Files
promptkit/internal/observation/recorder.go

319 lines
10 KiB
Go

// Package observation owns invocation-local execution identities, metadata
// accounting, and synchronous event delivery.
package observation
import (
"context"
"crypto/rand"
"fmt"
"math"
"time"
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
)
// Recorder accumulates one execution's measured attempts and usage. Its owner
// calls methods in execution order; it has no shared history or background work.
type Recorder struct {
runID string
startedAt time.Time
now func() time.Time
observer domain.ExecutionObserver
selection Selection
attempts uint64
retries uint64
repairCalls uint64
lastRepairGeneration uint64
finished bool
capacityWait time.Duration
generationDuration time.Duration
retryDelay time.Duration
usage usageTotals
}
// Selection contains resolved identities only. Unresolved fields stay empty.
type Selection struct {
PromptID string
PromptVersion string
ProfileID string
BackendID string
Model string
}
// AttemptData describes one completed active-permit acquisition and optional
// generation dispatch. Non-dispatched attempts contribute no token usage.
type AttemptData struct {
Generation uint64
Attempt uint64
Dispatched bool
CapacityWait time.Duration
GenerationDuration time.Duration
Outcome domain.ExecutionOutcome
Failure domain.ExecutionFailureKind
StatusCode int
Usage domain.TokenUsage
UsagePresence *domain.TokenUsagePresence
}
// DelayData describes one completed or interrupted retry wait.
type DelayData struct {
Generation uint64
Attempt uint64
Elapsed time.Duration
Failure domain.ExecutionFailureKind
}
// FinishData describes the terminal outcome after resource release.
type FinishData struct {
Outcome domain.ExecutionOutcome
Failure domain.ExecutionFailureKind
StatusCode int
RetriesExhausted bool
}
// New creates an invocation-local recorder and run ID. It returns the recorder
// even if ID creation fails, allowing an error terminal event with an empty ID.
func New(observer domain.ExecutionObserver) (*Recorder, error) {
runID, err := NewRunID()
return newRecorder(runID, time.Now, observer), err
}
func newRecorder(runID string, now func() time.Time, observer domain.ExecutionObserver) *Recorder {
return &Recorder{runID: runID, startedAt: now(), now: now, observer: observer,
usage: newUsageTotals()}
}
// NewRunID creates the identity shared by observations and successful results.
func NewRunID() (string, error) {
var b [16]byte
if _, err := rand.Read(b[:]); err != nil {
return "", err
}
b[6] = (b[6] & 0x0f) | 0x40
b[8] = (b[8] & 0x3f) | 0x80
return fmt.Sprintf("%08x-%04x-%04x-%04x-%012x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:16]), nil
}
// RunID returns this invocation's identity.
func (r *Recorder) RunID() string { return r.runID }
// Start delivers the opening event with only identities already resolved.
func (r *Recorder) Start() {
event := r.event(domain.ExecutionStarted)
r.deliver(event)
}
// UpdateSelection replaces the known resolved identities for later events.
func (r *Recorder) UpdateSelection(selection Selection) { r.selection = selection }
// Attempt records one completed acquisition and optional dispatch.
func (r *Recorder) Attempt(data AttemptData) {
r.capacityWait = addDuration(r.capacityWait, data.CapacityWait)
if data.Dispatched {
r.attempts++
if data.Attempt > 1 {
r.retries++
}
if data.Generation > 1 && data.Generation != r.lastRepairGeneration {
r.repairCalls++
r.lastRepairGeneration = data.Generation
}
r.generationDuration = addDuration(r.generationDuration, data.GenerationDuration)
r.usage.add(data.Usage, data.UsagePresence)
}
event := r.event(domain.ExecutionAttempt)
event.Generation = data.Generation
event.Attempt = data.Attempt
event.Dispatched = data.Dispatched
event.CapacityWait = data.CapacityWait
if data.Dispatched {
event.GenerationDuration = data.GenerationDuration
}
event.Outcome = data.Outcome
event.Failure = data.Failure
event.StatusCode = data.StatusCode
if data.Dispatched {
event.Usage = observedUsage(data.Usage, data.UsagePresence)
}
r.deliver(event)
}
// Delay records the measured time spent waiting after one failed attempt.
func (r *Recorder) Delay(data DelayData) {
r.retryDelay = addDuration(r.retryDelay, data.Elapsed)
event := r.event(domain.ExecutionRetryDelay)
event.Generation = data.Generation
event.Attempt = data.Attempt
event.RetryDelay = data.Elapsed
event.Failure = data.Failure
r.deliver(event)
}
// Finish delivers terminal totals once on a normal return. Its caller owns the
// operation's resource cleanup and failure classification.
func (r *Recorder) Finish(data FinishData) {
if r.finished {
return
}
r.finished = true
event := r.event(domain.ExecutionFinished)
event.Outcome = data.Outcome
event.Failure = data.Failure
event.StatusCode = data.StatusCode
event.RetriesExhausted = data.RetriesExhausted
event.Attempts = r.attempts
event.TransientRetries = r.retries
event.RepairCalls = r.repairCalls
event.CapacityWait = r.capacityWait
event.GenerationDuration = r.generationDuration
event.RetryDelay = r.retryDelay
event.Duration = event.Time.Sub(r.startedAt)
if event.Duration < 0 {
event.Duration = 0
}
event.Usage = r.usage.snapshot(r.attempts)
r.deliver(event)
}
func (r *Recorder) event(kind domain.ExecutionEventKind) domain.ExecutionEvent {
return domain.ExecutionEvent{
Kind: kind, RunID: r.runID, Time: r.now(), PromptID: r.selection.PromptID,
PromptVersion: r.selection.PromptVersion, ProfileID: r.selection.ProfileID,
BackendID: r.selection.BackendID, Model: r.selection.Model,
}
}
func (r *Recorder) deliver(event domain.ExecutionEvent) {
if r.observer != nil {
r.observer(event)
}
}
type contextKey struct{}
// WithoutExecutionState starts a public operation without inheriting another
// invocation's recorder or prepared-claim sink. Other context values and all
// cancellation/deadline behavior remain with the caller's context.
func WithoutExecutionState(ctx context.Context) context.Context {
return executionBoundary{Context: ctx}
}
type executionBoundary struct{ context.Context }
func (c executionBoundary) Value(key any) any {
switch key.(type) {
case contextKey, claimSinkKey:
return nil
default:
return c.Context.Value(key)
}
}
// WithRecorder carries one invocation's recorder through internal client calls.
func WithRecorder(ctx context.Context, recorder *Recorder) context.Context {
return context.WithValue(ctx, contextKey{}, recorder)
}
// FromContext returns the invocation-local recorder, if one was installed.
func FromContext(ctx context.Context) *Recorder {
recorder, _ := ctx.Value(contextKey{}).(*Recorder)
return recorder
}
type claimSinkKey struct{}
// WithClaimSink lets the facade learn which invocation actually claimed a
// prepared handle, so it can emit the terminal event after public mapping.
func WithClaimSink(ctx context.Context, sink **Recorder) context.Context {
return context.WithValue(ctx, claimSinkKey{}, sink)
}
// RegisterClaim records an invocation after its prepared handle is claimed.
func RegisterClaim(ctx context.Context, recorder *Recorder) {
if sink, ok := ctx.Value(claimSinkKey{}).(**Recorder); ok && sink != nil {
*sink = recorder
}
}
type usageField struct {
value int
known, complete bool
}
type usageTotals struct {
prompt, completion, total, cached, cacheWrite usageField
}
func newUsageTotals() usageTotals {
return usageTotals{prompt: usageField{complete: true}, completion: usageField{complete: true},
total: usageField{complete: true}, cached: usageField{complete: true}, cacheWrite: usageField{complete: true}}
}
func (u *usageTotals) add(usage domain.TokenUsage, presence *domain.TokenUsagePresence) {
known := usagePresence(usage, presence)
u.prompt.add(usage.PromptTokens, known.PromptTokens)
u.completion.add(usage.CompletionTokens, known.CompletionTokens)
u.total.add(usage.TotalTokens, known.TotalTokens)
u.cached.add(usage.CachedTokens, known.CachedTokens)
u.cacheWrite.add(usage.CacheWriteTokens, known.CacheWriteTokens)
}
func (u usageTotals) snapshot(attempts uint64) domain.ObservedTokenUsage {
return domain.ObservedTokenUsage{
PromptTokens: u.prompt.snapshot(attempts), CompletionTokens: u.completion.snapshot(attempts),
TotalTokens: u.total.snapshot(attempts), CachedTokens: u.cached.snapshot(attempts),
CacheWriteTokens: u.cacheWrite.snapshot(attempts),
}
}
func (f *usageField) add(value int, known bool) {
if known {
f.value += value
f.known = true
} else {
f.complete = false
}
}
func (f usageField) snapshot(attempts uint64) domain.ObservedTokenCount {
if attempts == 0 {
return domain.ObservedTokenCount{Known: true, Complete: true}
}
return domain.ObservedTokenCount{Value: f.value, Known: f.known, Complete: f.complete}
}
func observedUsage(usage domain.TokenUsage, presence *domain.TokenUsagePresence) domain.ObservedTokenUsage {
known := usagePresence(usage, presence)
return domain.ObservedTokenUsage{
PromptTokens: observedCount(usage.PromptTokens, known.PromptTokens),
CompletionTokens: observedCount(usage.CompletionTokens, known.CompletionTokens),
TotalTokens: observedCount(usage.TotalTokens, known.TotalTokens),
CachedTokens: observedCount(usage.CachedTokens, known.CachedTokens),
CacheWriteTokens: observedCount(usage.CacheWriteTokens, known.CacheWriteTokens),
}
}
func usagePresence(usage domain.TokenUsage, presence *domain.TokenUsagePresence) domain.TokenUsagePresence {
if presence != nil {
return *presence
}
return domain.TokenUsagePresence{
PromptTokens: usage.PromptTokens != 0, CompletionTokens: usage.CompletionTokens != 0,
TotalTokens: usage.TotalTokens != 0, CachedTokens: usage.CachedTokens != 0,
CacheWriteTokens: usage.CacheWriteTokens != 0,
}
}
func observedCount(value int, known bool) domain.ObservedTokenCount {
if !known {
return domain.ObservedTokenCount{}
}
return domain.ObservedTokenCount{Value: value, Known: true, Complete: true}
}
func addDuration(total, next time.Duration) time.Duration {
if next > 0 && total > time.Duration(math.MaxInt64)-next {
return time.Duration(math.MaxInt64)
}
return total + next
}