Verify retry accounting and execution summary propagation
This commit is contained in:
245
internal/framework/llm/promptkit_observation_test.go
Normal file
245
internal/framework/llm/promptkit_observation_test.go
Normal file
@@ -0,0 +1,245 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||
"gitea.maximumdirect.net/eric/promptkit"
|
||||
)
|
||||
|
||||
func newObservedHTTPClient(t *testing.T, handler http.HandlerFunc, retries int, transport http.RoundTripper) *PromptKitClient {
|
||||
t.Helper()
|
||||
server := httptest.NewServer(handler)
|
||||
t.Cleanup(server.Close)
|
||||
client, err := NewPromptKitClient(PromptKitClientConfig{
|
||||
Assets: newTestPromptKitAssets(t), TransientRetries: &retries,
|
||||
HTTPClient: &http.Client{Transport: transport},
|
||||
EngineOptions: []promptkit.Option{promptkit.WithProfiles(promptkit.OpenAICompatibleProfile(promptkit.OpenAICompatibleProfileConfig{
|
||||
ID: "default-profile", Endpoint: server.URL + "/v1", Model: "test-model",
|
||||
}))},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return client
|
||||
}
|
||||
|
||||
func observedRequest() contracts.StructuredCompletionRequest {
|
||||
return contracts.StructuredCompletionRequest{
|
||||
PromptID: "adapter.test", SessionID: "observation-test",
|
||||
Inputs: contracts.LLMInputSet{"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{}`), "", "")},
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromptKitObservationsComposeTransientRetriesAndRepair(t *testing.T) {
|
||||
var calls atomic.Int32
|
||||
client := newObservedHTTPClient(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
switch calls.Add(1) {
|
||||
case 1, 3:
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
_, _ = io.WriteString(w, `{"error":{"message":"try again"}}`)
|
||||
case 2:
|
||||
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"bad\":true}"}}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5,"prompt_tokens_details":{"cached_tokens":0}}}`)
|
||||
case 4:
|
||||
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"ok\":true}"}}],"usage":{"prompt_tokens":7,"completion_tokens":0,"total_tokens":7}}`)
|
||||
default:
|
||||
t.Error("unexpected extra HTTP attempt")
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
}
|
||||
}, 1, nil)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
req := observedRequest()
|
||||
repairs := 1
|
||||
req.StructuredOutputRepairAttempts = &repairs
|
||||
var out struct{ OK bool }
|
||||
response, err := client.CompleteStructured(ctx, req, &out)
|
||||
if err != nil || !out.OK {
|
||||
t.Fatalf("completion = (%+v, %v), output = %+v", response, err, out)
|
||||
}
|
||||
summary := response.Execution
|
||||
if calls.Load() != 4 || summary == nil || summary.Attempts != 4 || summary.TransientRetries != 2 || summary.RepairCalls != 1 || summary.Outcome != "success" || summary.RetriesExhausted {
|
||||
t.Fatalf("calls = %d, execution = %+v", calls.Load(), summary)
|
||||
}
|
||||
wantUsage := contracts.LLMObservedTokenUsage{
|
||||
PromptTokens: contracts.LLMObservedTokenCount{Value: 10, Known: true},
|
||||
CompletionTokens: contracts.LLMObservedTokenCount{Value: 2, Known: true},
|
||||
TotalTokens: contracts.LLMObservedTokenCount{Value: 12, Known: true},
|
||||
CachedTokens: contracts.LLMObservedTokenCount{Known: true},
|
||||
}
|
||||
// Failed HTTP attempts contribute unknown usage, not known zero. Neither
|
||||
// terminal totals nor legacy result totals should be added a second time.
|
||||
if summary.Usage != wantUsage {
|
||||
t.Fatalf("observed usage = %+v, want %+v", summary.Usage, wantUsage)
|
||||
}
|
||||
if response.PromptTokens != 10 || response.CompletionTokens != 2 || response.TotalTokens != 12 || response.RepairAttempts != 1 {
|
||||
t.Fatalf("legacy result accounting changed: %+v", response)
|
||||
}
|
||||
}
|
||||
|
||||
type closeSignalingBody struct {
|
||||
io.ReadCloser
|
||||
closed chan<- struct{}
|
||||
}
|
||||
|
||||
func (b closeSignalingBody) Close() error {
|
||||
err := b.ReadCloser.Close()
|
||||
b.closed <- struct{}{}
|
||||
return err
|
||||
}
|
||||
|
||||
func TestPromptKitObservationSurvivesCanceledRetryWait(t *testing.T) {
|
||||
consumed := make(chan struct{}, 2)
|
||||
transport := http.DefaultTransport.(*http.Transport).Clone()
|
||||
t.Cleanup(transport.CloseIdleConnections)
|
||||
var calls atomic.Int32
|
||||
client := newObservedHTTPClient(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
calls.Add(1)
|
||||
w.Header().Set("Retry-After", "3600")
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
_, _ = io.WriteString(w, `{"error":{"message":"busy"}}`)
|
||||
}, 1, roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
response, err := transport.RoundTrip(req)
|
||||
if err == nil {
|
||||
response.Body = closeSignalingBody{ReadCloser: response.Body, closed: consumed}
|
||||
}
|
||||
return response, err
|
||||
}))
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
type result struct {
|
||||
response contracts.StructuredCompletionResponse
|
||||
err error
|
||||
}
|
||||
done := make(chan result, 1)
|
||||
go func() {
|
||||
var out map[string]any
|
||||
response, err := client.CompleteStructured(ctx, observedRequest(), &out)
|
||||
done <- result{response, err}
|
||||
}()
|
||||
// The failure has been consumed, so cancellation must interrupt the pending
|
||||
// retry wait rather than depend on another provider response or its timeout.
|
||||
select {
|
||||
case <-consumed:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("provider response was not consumed")
|
||||
}
|
||||
cancel()
|
||||
select {
|
||||
case got := <-done:
|
||||
if !errors.Is(got.err, context.Canceled) || got.response.Execution == nil {
|
||||
t.Fatalf("canceled completion = (%+v, %v)", got.response, got.err)
|
||||
}
|
||||
summary := got.response.Execution
|
||||
if summary.Outcome != "error" || summary.FailureKind != "canceled" || summary.Attempts != 1 || summary.RetriesExhausted || calls.Load() != 1 {
|
||||
t.Fatalf("canceled execution = %+v, HTTP calls = %d", summary, calls.Load())
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("cancellation did not interrupt retry waiting")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromptKitObservationsSeparatePreparationAndExecutionFailure(t *testing.T) {
|
||||
var calls atomic.Int32
|
||||
client := newObservedHTTPClient(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
calls.Add(1)
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
_, _ = io.WriteString(w, `{"error":{"message":"busy"}}`)
|
||||
}, 1, nil)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
req := observedRequest()
|
||||
req.PromptID = "missing-prompt"
|
||||
var out map[string]any
|
||||
response, err := client.CompleteStructured(ctx, req, &out)
|
||||
if err == nil || response.Execution != nil || calls.Load() != 0 {
|
||||
t.Fatalf("preparation failure = (%+v, %v), calls = %d", response, err, calls.Load())
|
||||
}
|
||||
response, err = client.CompleteStructured(ctx, observedRequest(), &out)
|
||||
var generation *contracts.LLMGenerationError
|
||||
if !errors.As(err, &generation) || response.Execution == nil {
|
||||
t.Fatalf("execution failure = (%+v, %v)", response, err)
|
||||
}
|
||||
summary := response.Execution
|
||||
if summary.Outcome != "error" || summary.FailureKind != "provider_status" || summary.StatusCode != http.StatusServiceUnavailable || !summary.RetriesExhausted || summary.Attempts != 2 || summary.TransientRetries != 1 || calls.Load() != 2 {
|
||||
t.Fatalf("failed execution = %+v, calls = %d", summary, calls.Load())
|
||||
}
|
||||
if summary.Usage != (contracts.LLMObservedTokenUsage{}) {
|
||||
t.Fatalf("failed attempts should have unknown usage: %+v", summary.Usage)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromptKitObservationsAreIsolatedAcrossConcurrentCalls(t *testing.T) {
|
||||
arrived := make(chan struct{}, 2)
|
||||
release := make(chan struct{})
|
||||
client := newObservedHTTPClient(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
body, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
tokens := 11
|
||||
if strings.Contains(string(body), "second-call") {
|
||||
tokens = 22
|
||||
}
|
||||
arrived <- struct{}{}
|
||||
select {
|
||||
case <-release:
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"choices": []any{map[string]any{"message": map[string]string{"content": `{"ok":true}`}}},
|
||||
"usage": map[string]int{"prompt_tokens": tokens, "completion_tokens": 0, "total_tokens": tokens},
|
||||
})
|
||||
}, 0, nil)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
type result struct {
|
||||
want int
|
||||
response contracts.StructuredCompletionResponse
|
||||
err error
|
||||
}
|
||||
done := make(chan result, 2)
|
||||
for name, tokens := range map[string]int{"first-call": 11, "second-call": 22} {
|
||||
go func() {
|
||||
req := observedRequest()
|
||||
req.Vars = map[string]any{"custom": name}
|
||||
var out map[string]any
|
||||
response, err := client.CompleteStructured(ctx, req, &out)
|
||||
done <- result{tokens, response, err}
|
||||
}()
|
||||
}
|
||||
for i := 0; i < 2; i++ {
|
||||
select {
|
||||
case <-arrived:
|
||||
case <-ctx.Done():
|
||||
t.Fatal("requests did not overlap")
|
||||
}
|
||||
}
|
||||
close(release)
|
||||
for i := 0; i < 2; i++ {
|
||||
select {
|
||||
case got := <-done:
|
||||
summary := got.response.Execution
|
||||
if got.err != nil || summary == nil {
|
||||
t.Fatalf("concurrent completion = (%+v, %v)", got.response, got.err)
|
||||
}
|
||||
if summary.Attempts != 1 || summary.Usage.PromptTokens != (contracts.LLMObservedTokenCount{Value: got.want, Known: true, Complete: true}) || summary.Usage.CompletionTokens != (contracts.LLMObservedTokenCount{Known: true, Complete: true}) || summary.Usage.CachedTokens != (contracts.LLMObservedTokenCount{}) {
|
||||
t.Fatalf("execution for %d tokens = %+v", got.want, summary)
|
||||
}
|
||||
case <-ctx.Done():
|
||||
t.Fatal("concurrent completions did not return")
|
||||
}
|
||||
}
|
||||
}
|
||||
122
internal/framework/pipeline/debug_execution_test.go
Normal file
122
internal/framework/pipeline/debug_execution_test.go
Normal file
@@ -0,0 +1,122 @@
|
||||
package pipeline
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
|
||||
)
|
||||
|
||||
type executionDebugClient struct {
|
||||
response contracts.StructuredCompletionResponse
|
||||
err error
|
||||
}
|
||||
|
||||
func (c executionDebugClient) CompleteStructured(context.Context, contracts.StructuredCompletionRequest, any) (contracts.StructuredCompletionResponse, error) {
|
||||
return c.response, c.err
|
||||
}
|
||||
|
||||
func TestScheduledDebugClientPreservesExecutionThroughTerminalRecords(t *testing.T) {
|
||||
providerErr := contracts.NewLLMGenerationError(503, "provider unavailable")
|
||||
const content = `{"ok":true}`
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
err error
|
||||
content string
|
||||
debugResponse bool
|
||||
}{
|
||||
{name: "provider failure without body", err: providerErr},
|
||||
{name: "canceled without body", err: context.Canceled},
|
||||
{name: "split debug response", content: content, debugResponse: true},
|
||||
{name: "split fallback response", content: content},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
want := contracts.LLMExecutionSummary{
|
||||
Attempts: 2, TransientRetries: 1, Outcome: "success",
|
||||
CapacityWaitMS: 1, GenerationMS: 3, RetryDelayMS: 5, DurationMS: 9,
|
||||
Usage: contracts.LLMObservedTokenUsage{
|
||||
PromptTokens: contracts.LLMObservedTokenCount{Value: 7, Known: true},
|
||||
CompletionTokens: contracts.LLMObservedTokenCount{Known: true},
|
||||
},
|
||||
}
|
||||
if test.err != nil {
|
||||
want.Outcome = "error"
|
||||
want.FailureKind = "provider_status"
|
||||
want.StatusCode = 503
|
||||
want.RetriesExhausted = true
|
||||
if errors.Is(test.err, context.Canceled) {
|
||||
want.FailureKind = "canceled"
|
||||
want.StatusCode = 0
|
||||
want.RetriesExhausted = false
|
||||
}
|
||||
}
|
||||
execution := want
|
||||
response := contracts.StructuredCompletionResponse{Execution: &execution}
|
||||
if test.debugResponse {
|
||||
response.Debug = &contracts.LLMDebugMaterial{Response: &contracts.LLMDebugResponse{Content: test.content}}
|
||||
} else {
|
||||
response.Content = json.RawMessage(test.content)
|
||||
}
|
||||
scheduler, err := llm.NewScheduler(1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
recorder := newCapturedDebugRecorder()
|
||||
client := WithDebugLLMRecording(llm.NewScheduledClient(executionDebugClient{response, test.err}, scheduler), recorder)
|
||||
ctx, scope := withDebugLLMScope(context.Background(), "attempt")
|
||||
got, err := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{StageName: "extract"}, &struct{}{})
|
||||
if !errors.Is(err, test.err) || got.Execution == nil || *got.Execution != want {
|
||||
t.Fatalf("scheduled debug completion = (%+v, %v), want summary and original error", got, err)
|
||||
}
|
||||
terminal := newAttemptTerminalRecorder(recorder, "attempt/terminal", "extract", scope, debugTimedEnvelope{Stage: "extract"})
|
||||
if err := terminal.record(nil, test.err); !errors.Is(err, test.err) {
|
||||
t.Fatalf("terminal recording error = %v", err)
|
||||
}
|
||||
envelope := recorder.envelope(t, "attempt/response-0001.json")
|
||||
encoded, err := json.Marshal(envelope.Payload)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var artifact debugLLMResponseArtifact
|
||||
if err := json.Unmarshal(encoded, &artifact); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if artifact.Execution == nil || *artifact.Execution != want {
|
||||
t.Fatalf("persisted response lost execution summary: %+v", artifact)
|
||||
}
|
||||
if !test.debugResponse && (artifact.Fallback == nil || artifact.Fallback.Execution == nil || *artifact.Fallback.Execution != want) {
|
||||
t.Fatalf("persisted fallback lost execution summary: %+v", artifact.Fallback)
|
||||
}
|
||||
if test.content == "" {
|
||||
if artifact.ContentPath != "" {
|
||||
t.Fatalf("bodyless failure has content path %q", artifact.ContentPath)
|
||||
}
|
||||
} else {
|
||||
var compact bytes.Buffer
|
||||
if artifact.ContentPath == "" || json.Compact(&compact, recorder.bytes[artifact.ContentPath]) != nil || compact.String() != test.content {
|
||||
t.Fatalf("split content not preserved at %q", artifact.ContentPath)
|
||||
}
|
||||
if (artifact.Fallback != nil && artifact.Fallback.Content != "") || (artifact.Response != nil && artifact.Response.Content != "") {
|
||||
t.Fatal("split content remains in response metadata")
|
||||
}
|
||||
}
|
||||
refs := recorder.envelope(t, "attempt/terminal.json").LLMCalls
|
||||
if len(refs) != 1 || refs[0].Execution == nil || *refs[0].Execution != want || refs[0].Error != (test.err != nil) {
|
||||
t.Fatalf("terminal references = %+v", refs)
|
||||
}
|
||||
// Decode the persisted JSON, not just Go defaults: omitted false and
|
||||
// zero fields would erase the presence-aware usage contract.
|
||||
usage := envelope.Payload.(map[string]any)["execution"].(map[string]any)["usage"].(map[string]any)
|
||||
unknown := usage["cached_tokens"].(map[string]any)
|
||||
for _, key := range []string{"value", "known", "complete"} {
|
||||
if _, ok := unknown[key]; !ok {
|
||||
t.Fatalf("unknown token usage omits %q: %+v", key, unknown)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user