eric/audita
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 6 Wiki Activity

Collect PromptKit execution observations with logical attribution

Browse Source
This commit is contained in:
Eric Rakestraw
2026-09-24 02:23:57 +00:00
parent 4812032c9e
commit 4bfc5f3456
11 changed files with 457 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
internal/framework/llm/adapter.go
View File

@@ -16,8 +16,9 @@ import (
// CallHooks receive one logical start and finish record and the execution
// events for this call. Call records can contain transcript content.
type CallHooks struct {
OnCall func(structuredoutput.CallRecord)
OnEvent pk.ExecutionObserver
OnCall func(structuredoutput.CallRecord)
OnEvent pk.ExecutionObserver
OnEventWithCall func(uint64, structuredoutput.DiagnosticContext, pk.ExecutionEvent)
}
type RoleClient struct {
@@ -38,7 +39,7 @@ func (r *Runtime) NewRoleClient(role string, hooks CallHooks) (*RoleClient, erro
func (c *RoleClient) CompleteStructured(parent context.Context, req structuredoutput.Request, out any) (response structuredoutput.Response, err error) {
ctx, cancel := context.WithTimeout(parent, c.runtime.Config.LLM.OperationTimeout)
defer cancel()
record := structuredoutput.CallRecord{Context: req.Context, Started: time.Now().UTC()}
record := structuredoutput.CallRecord{CallID: c.runtime.nextCallID.Add(1), Context: copyContext(req.Context), Started: time.Now().UTC()}
record.Context.Role = c.role
if c.hooks.OnCall != nil {
c.hooks.OnCall(record)
@@ -78,7 +79,14 @@ func (c *RoleClient) CompleteStructured(parent context.Context, req structuredou
PromptID: req.PromptID, PromptVersion: req.PromptVersion,
ProfileID: profile.ProfileID, APIKey: apiKey, Vars: vars,
Validation: &pk.OutputContract{Format: pk.FormatJSON, ValidationMode: pk.ValidationJSONSchema, SchemaPath: schemaPath, RepairAttempts: c.runtime.Config.LLM.RepairAttempts},
Observer: c.hooks.OnEvent,
Observer: func(event pk.ExecutionEvent) {
if c.hooks.OnEvent != nil {
c.hooks.OnEvent(event)
}
if c.hooks.OnEventWithCall != nil {
c.hooks.OnEventWithCall(record.CallID, record.Context, event)
}
},
})
if prepErr != nil {
return response, &structuredoutput.Failure{Kind: structuredoutput.FailureRuntime, Message: "could not prepare structured model call", Cause: prepErr}

209
internal/framework/llm/observations.go Normal file
View File

@@ -0,0 +1,209 @@
package llm
import (
"sort"
"sync"
"time"
"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
pk "gitea.maximumdirect.net/eric/promptkit"
)
// TokenCount retains the distinction between zero and unavailable usage.
type TokenCount struct {
Value int `json:"value"`
Known bool `json:"known"`
Complete bool `json:"complete"`
}
type Usage struct {
Prompt TokenCount `json:"prompt"`
Completion TokenCount `json:"completion"`
Total TokenCount `json:"total"`
Cached TokenCount `json:"cached"`
CacheWrite TokenCount `json:"cache_write"`
}
// ExecutionObservation contains execution metadata only. It never contains
// endpoint URLs, credentials, messages, provider errors, or generated text.
type ExecutionObservation struct {
CallID uint64 `json:"call_id"`
Context structuredoutput.DiagnosticContext `json:"context"`
Kind string `json:"kind"`
RunID string `json:"run_id,omitempty"`
Time time.Time `json:"time"`
BackendID string `json:"backend_id,omitempty"`
ProfileID string `json:"profile_id,omitempty"`
Model string `json:"model,omitempty"`
Generation uint64 `json:"generation,omitempty"`
Attempt uint64 `json:"attempt,omitempty"`
Dispatched bool `json:"dispatched"`
CapacityWait time.Duration `json:"capacity_wait_ns"`
GenerationDuration time.Duration `json:"generation_duration_ns"`
RetryDelay time.Duration `json:"retry_delay_ns"`
Duration time.Duration `json:"duration_ns"`
Outcome string `json:"outcome,omitempty"`
Failure string `json:"failure,omitempty"`
StatusCode int `json:"status_code,omitempty"`
RetriesExhausted bool `json:"retries_exhausted,omitempty"`
Attempts uint64 `json:"attempts"`
TransientRetries uint64 `json:"transient_retries"`
StructuralRepairs uint64 `json:"structural_repairs"`
Usage Usage `json:"usage"`
}
type LogicalCall struct {
CallID uint64 `json:"call_id"`
Context structuredoutput.DiagnosticContext `json:"context"`
Started time.Time `json:"started"`
Finished time.Time `json:"finished,omitempty"`
WallDuration time.Duration `json:"wall_duration_ns"`
Failure structuredoutput.FailureKind `json:"failure,omitempty"`
}
type ObservationSnapshot struct {
Calls []LogicalCall `json:"calls"`
Events []ExecutionObservation `json:"events"`
}
type ExecutionSummary struct {
LogicalCalls uint64 `json:"logical_calls"`
ProviderAttempts uint64 `json:"provider_attempts"`
TransientRetries uint64 `json:"transient_retries"`
StructuralRepairs uint64 `json:"structural_repairs"`
Usage Usage `json:"usage"`
}
// ObservationCollector accepts short synchronous callbacks from parallel calls.
// Snapshot should be taken after the runner has joined active section work.
type ObservationCollector struct {
mu sync.Mutex
calls map[uint64]LogicalCall
events []ExecutionObservation
}
func NewObservationCollector() *ObservationCollector {
return &ObservationCollector{calls: make(map[uint64]LogicalCall)}
}
func (c *ObservationCollector) Hooks() CallHooks {
return CallHooks{OnCall: c.recordCall, OnEventWithCall: c.recordEvent}
}
func copyContext(value structuredoutput.DiagnosticContext) structuredoutput.DiagnosticContext {
if value.SectionIndex != nil {
index := *value.SectionIndex
value.SectionIndex = &index
}
if value.BatchIndex != nil {
index := *value.BatchIndex
value.BatchIndex = &index
}
return value
}
func (c *ObservationCollector) recordCall(record structuredoutput.CallRecord) {
if c == nil {
return
}
c.mu.Lock()
defer c.mu.Unlock()
call := LogicalCall{CallID: record.CallID, Context: copyContext(record.Context), Started: record.Started, Finished: record.Finished, Failure: record.Failure}
if !record.Finished.IsZero() {
call.WallDuration = record.Finished.Sub(record.Started)
}
c.calls[record.CallID] = call
}
func (c *ObservationCollector) recordEvent(callID uint64, context structuredoutput.DiagnosticContext, event pk.ExecutionEvent) {
if c == nil {
return
}
count := func(value pk.ObservedTokenCount) TokenCount {
return TokenCount{Value: value.Value, Known: value.Known, Complete: value.Complete}
}
observation := ExecutionObservation{
CallID: callID, Context: copyContext(context), Kind: string(event.Kind), RunID: event.RunID,
Time: event.Time, BackendID: event.BackendID, ProfileID: event.ProfileID, Model: event.Model,
Generation: event.Generation, Attempt: event.Attempt, Dispatched: event.Dispatched,
CapacityWait: event.CapacityWait, GenerationDuration: event.GenerationDuration, RetryDelay: event.RetryDelay, Duration: event.Duration,
Outcome: string(event.Outcome), Failure: string(event.Failure), StatusCode: event.StatusCode, RetriesExhausted: event.RetriesExhausted,
Attempts: event.Attempts, TransientRetries: event.TransientRetries, StructuralRepairs: event.RepairCalls,
Usage: Usage{Prompt: count(event.Usage.PromptTokens), Completion: count(event.Usage.CompletionTokens), Total: count(event.Usage.TotalTokens), Cached: count(event.Usage.CachedTokens), CacheWrite: count(event.Usage.CacheWriteTokens)},
}
c.mu.Lock()
c.events = append(c.events, observation)
c.mu.Unlock()
}
func (c *ObservationCollector) Snapshot() ObservationSnapshot {
if c == nil {
return ObservationSnapshot{}
}
c.mu.Lock()
defer c.mu.Unlock()
snapshot := ObservationSnapshot{Calls: make([]LogicalCall, 0, len(c.calls)), Events: make([]ExecutionObservation, len(c.events))}
for _, call := range c.calls {
call.Context = copyContext(call.Context)
snapshot.Calls = append(snapshot.Calls, call)
}
copy(snapshot.Events, c.events)
for i := range snapshot.Events {
snapshot.Events[i].Context = copyContext(snapshot.Events[i].Context)
}
sort.Slice(snapshot.Calls, func(i, j int) bool { return snapshot.Calls[i].CallID < snapshot.Calls[j].CallID })
sort.SliceStable(snapshot.Events, func(i, j int) bool {
if snapshot.Events[i].CallID != snapshot.Events[j].CallID {
return snapshot.Events[i].CallID < snapshot.Events[j].CallID
}
return snapshot.Events[i].Time.Before(snapshot.Events[j].Time)
})
return snapshot
}
// Summarize counts terminal usage once per logical call. PromptKit's terminal
// usage already aggregates its attempts, retries, and repairs.
func Summarize(snapshot ObservationSnapshot) ExecutionSummary {
result := ExecutionSummary{LogicalCalls: uint64(len(snapshot.Calls))}
terminal := make(map[uint64]ExecutionObservation)
for _, event := range snapshot.Events {
if event.Kind == string(pk.ExecutionFinished) {
terminal[event.CallID] = event
}
}
dispatched := false
aggregate := func(dst *TokenCount, source TokenCount) {
if source.Known {
dst.Value += source.Value
dst.Known = true
}
dst.Complete = dst.Complete && source.Complete
}
result.Usage = Usage{Prompt: TokenCount{Complete: true}, Completion: TokenCount{Complete: true}, Total: TokenCount{Complete: true}, Cached: TokenCount{Complete: true}, CacheWrite: TokenCount{Complete: true}}
for _, call := range snapshot.Calls {
event, ok := terminal[call.CallID]
if !ok {
continue
}
result.ProviderAttempts += event.Attempts
result.TransientRetries += event.TransientRetries
result.StructuralRepairs += event.StructuralRepairs
if event.Attempts == 0 {
continue
}
dispatched = true
aggregate(&result.Usage.Prompt, event.Usage.Prompt)
aggregate(&result.Usage.Completion, event.Usage.Completion)
aggregate(&result.Usage.Total, event.Usage.Total)
aggregate(&result.Usage.Cached, event.Usage.Cached)
aggregate(&result.Usage.CacheWrite, event.Usage.CacheWrite)
}
if !dispatched {
result.Usage.Prompt.Known = true
result.Usage.Completion.Known = true
result.Usage.Total.Known = true
result.Usage.Cached.Known = true
result.Usage.CacheWrite.Known = true
}
return result
}

126
internal/framework/llm/observations_test.go Normal file
View File

@@ -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)
}
}

8
internal/framework/llm/runtime.go
View File

@@ -9,6 +9,7 @@ import (
"path/filepath"
"sort"
"strings"
"sync/atomic"
"gitea.maximumdirect.net/eric/audita/internal/core/config"
"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
@@ -19,9 +20,10 @@ import (
// Runtime owns one PromptKit engine and the validated selections for both roles.
// The same engine is shared by the role-bound clients during processing.
type Runtime struct {
Config config.Config
Engine *pk.Engine
Profiles map[string]pk.ProfileInspection
Config config.Config
Engine *pk.Engine
Profiles map[string]pk.ProfileInspection
nextCallID atomic.Uint64
}
// AssetError indicates a broken embedded prompt or schema rather than an