Collect PromptKit execution observations with logical attribution

2026-09-24 02:23:57 +00:00
parent 4812032c9e
commit 4bfc5f3456
11 changed files with 457 additions and 14 deletions
--- a/internal/framework/llm/observations_test.go
+++ b/internal/framework/llm/observations_test.go
@@ -0,0 +1,126 @@
+package llm
+
+import (
+	"context"
+	"encoding/json"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
+	"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
+	pk "gitea.maximumdirect.net/eric/promptkit"
+)
+
+func TestObservationSummarySeparatesLogicalCallsAttemptsAndUsage(t *testing.T) {
+	collector := NewObservationCollector()
+	hooks := collector.Hooks()
+	now := time.Now().UTC()
+	call := func(id uint64) {
+		hooks.OnCall(structuredoutput.CallRecord{CallID: id, Context: structuredoutput.DiagnosticContext{Role: "proposal"}, Started: now})
+		hooks.OnCall(structuredoutput.CallRecord{CallID: id, Context: structuredoutput.DiagnosticContext{Role: "proposal"}, Started: now, Finished: now.Add(time.Second)})
+	}
+	call(1) // preparation failed before PromptKit emitted events
+	call(2) // capacity rejected without a dispatched attempt
+	hooks.OnEventWithCall(2, structuredoutput.DiagnosticContext{Role: "proposal"}, pk.ExecutionEvent{Kind: pk.ExecutionAttempt, Dispatched: false})
+	hooks.OnEventWithCall(2, structuredoutput.DiagnosticContext{Role: "proposal"}, pk.ExecutionEvent{Kind: pk.ExecutionFinished, Attempts: 0})
+	summary := Summarize(collector.Snapshot())
+	if summary.LogicalCalls != 2 || summary.ProviderAttempts != 0 || summary.Usage.Total != (TokenCount{Known: true, Complete: true}) {
+		t.Fatalf("zero-attempt accounting: %+v", summary)
+	}
+
+	call(3)
+	context := structuredoutput.DiagnosticContext{Role: "validation", ModuleInstance: "grammar", ValidatorName: "editorial_review"}
+	hooks.OnEventWithCall(3, context, pk.ExecutionEvent{Kind: pk.ExecutionAttempt, Dispatched: true, Usage: pk.ObservedTokenUsage{TotalTokens: pk.ObservedTokenCount{Value: 100, Known: true, Complete: true}}})
+	hooks.OnEventWithCall(3, context, pk.ExecutionEvent{Kind: pk.ExecutionFinished, Attempts: 3, TransientRetries: 1, RepairCalls: 1, Usage: pk.ObservedTokenUsage{
+		PromptTokens: pk.ObservedTokenCount{Value: 5, Known: true, Complete: false},
+		TotalTokens:  pk.ObservedTokenCount{Value: 8, Known: true, Complete: false},
+	}})
+	call(4)
+	hooks.OnEventWithCall(4, context, pk.ExecutionEvent{Kind: pk.ExecutionFinished, Attempts: 1, Usage: pk.ObservedTokenUsage{PromptTokens: pk.ObservedTokenCount{Known: false, Complete: false}, TotalTokens: pk.ObservedTokenCount{Known: false, Complete: false}}})
+	snapshot := collector.Snapshot()
+	summary = Summarize(snapshot)
+	if summary.LogicalCalls != 4 || summary.ProviderAttempts != 4 || summary.TransientRetries != 1 || summary.StructuralRepairs != 1 || summary.Usage.Total.Value != 8 || !summary.Usage.Total.Known || summary.Usage.Total.Complete || summary.Usage.Prompt.Value != 5 || !summary.Usage.Prompt.Known || summary.Usage.Prompt.Complete || summary.Usage.Cached.Known || summary.Usage.Cached.Complete {
+		t.Fatalf("partial usage or double count: %+v", summary)
+	}
+	raw, err := json.Marshal(snapshot)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, forbidden := range []string{"https://", "Bearer ", "raw_output", "messages", "api_key"} {
+		if strings.Contains(string(raw), forbidden) {
+			t.Fatalf("observation leaked payload field %q", forbidden)
+		}
+	}
+}
+
+func TestObservationCollectorConcurrentAttributionAndSnapshotCopy(t *testing.T) {
+	collector := NewObservationCollector()
+	hooks := collector.Hooks()
+	const calls = 64
+	var wg sync.WaitGroup
+	for i := 1; i <= calls; i++ {
+		wg.Add(1)
+		go func(id int) {
+			defer wg.Done()
+			index := id
+			context := structuredoutput.DiagnosticContext{Role: "proposal", ModuleInstance: "grammar", SectionIndex: &index}
+			started := time.Now().UTC()
+			hooks.OnCall(structuredoutput.CallRecord{CallID: uint64(id), Context: context, Started: started})
+			hooks.OnEventWithCall(uint64(id), context, pk.ExecutionEvent{Kind: pk.ExecutionFinished, RunID: "run", Attempts: 1, Usage: pk.ObservedTokenUsage{TotalTokens: pk.ObservedTokenCount{Value: 1, Known: true, Complete: true}}})
+			hooks.OnCall(structuredoutput.CallRecord{CallID: uint64(id), Context: context, Started: started, Finished: time.Now().UTC()})
+			index = -1
+		}(i)
+	}
+	wg.Wait()
+	snapshot := collector.Snapshot()
+	if len(snapshot.Calls) != calls || len(snapshot.Events) != calls || Summarize(snapshot).ProviderAttempts != calls {
+		t.Fatalf("parallel records: calls=%d events=%d", len(snapshot.Calls), len(snapshot.Events))
+	}
+	for _, call := range snapshot.Calls {
+		if call.Context.SectionIndex == nil || *call.Context.SectionIndex != int(call.CallID) {
+			t.Fatalf("attribution lost: %+v", call)
+		}
+	}
+	*snapshot.Calls[0].Context.SectionIndex = -2
+	again := collector.Snapshot()
+	if *again.Calls[0].Context.SectionIndex == -2 {
+		t.Fatal("snapshot aliases collector context")
+	}
+}
+
+func TestRoleClientFeedsCollectorAcrossPreparationFailure(t *testing.T) {
+	runtime, count, _ := adapterFixture(t, []reply{{status: 200, content: `{"corrections":[]}`}}, 0, 0, nil)
+	collector := NewObservationCollector()
+	client, err := runtime.NewRoleClient("proposal", collector.Hooks())
+	if err != nil {
+		t.Fatal(err)
+	}
+	request := adapterRequest(responseschema.CorrectionSetKey)
+	request.Context.ModuleInstance = "glossary_1"
+	index := 3
+	request.Context.SectionIndex = &index
+	missing := request
+	missing.PromptID = "missing.prompt"
+	var output map[string]any
+	if _, err := client.CompleteStructured(context.Background(), missing, &output); err == nil {
+		t.Fatal("preparation error missing")
+	}
+	if _, err := client.CompleteStructured(context.Background(), request, &output); err != nil {
+		t.Fatal(err)
+	}
+	snapshot := collector.Snapshot()
+	summary := Summarize(snapshot)
+	if count.Load() != 1 || summary.LogicalCalls != 2 || summary.ProviderAttempts != 1 || len(snapshot.Events) == 0 {
+		t.Fatalf("collector accounting: %+v", summary)
+	}
+	for _, event := range snapshot.Events {
+		if event.Context.ModuleInstance != "glossary_1" || event.Context.SectionIndex == nil || *event.Context.SectionIndex != 3 || event.Context.Role != "proposal" {
+			t.Fatalf("event attribution: %+v", event)
+		}
+	}
+	if snapshot.Calls[0].Failure != structuredoutput.FailureRuntime || snapshot.Calls[1].Failure != "" {
+		t.Fatalf("logical failures: %+v", snapshot.Calls)
+	}
+}