Collect PromptKit execution observations with logical attribution
This commit is contained in:
209
internal/framework/llm/observations.go
Normal file
209
internal/framework/llm/observations.go
Normal 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
|
||||
}
|
||||
Reference in New Issue
Block a user