246 lines
9.0 KiB
Go
246 lines
9.0 KiB
Go
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")
|
|
}
|
|
}
|
|
}
|