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")
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user