80 lines
3.1 KiB
Go
80 lines
3.1 KiB
Go
package main
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
|
|
pk "gitea.maximumdirect.net/eric/promptkit"
|
|
)
|
|
|
|
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"`
|
|
}
|
|
type Observation struct {
|
|
Kind string `json:"kind"`
|
|
RunID string `json:"run_id"`
|
|
Role string `json:"role"`
|
|
Backend string `json:"backend"`
|
|
Profile string `json:"profile"`
|
|
Generation uint64 `json:"generation,omitempty"`
|
|
Attempt uint64 `json:"attempt,omitempty"`
|
|
Dispatched bool `json:"dispatched,omitempty"`
|
|
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"`
|
|
Failure string `json:"failure,omitempty"`
|
|
StatusCode int `json:"status_code,omitempty"`
|
|
RetriesExhausted bool `json:"retries_exhausted,omitempty"`
|
|
Attempts uint64 `json:"attempts,omitempty"`
|
|
Retries uint64 `json:"retries,omitempty"`
|
|
Repairs uint64 `json:"repairs,omitempty"`
|
|
Usage Usage `json:"usage"`
|
|
}
|
|
|
|
type Recorder struct {
|
|
mu sync.Mutex
|
|
records []Observation
|
|
}
|
|
|
|
func (r *Recorder) observer(role string) pk.ExecutionObserver {
|
|
return func(e pk.ExecutionEvent) {
|
|
count := func(c pk.ObservedTokenCount) TokenCount { return TokenCount{c.Value, c.Known, c.Complete} }
|
|
o := Observation{Kind: string(e.Kind), RunID: e.RunID, Role: role, Backend: e.BackendID, Profile: e.ProfileID,
|
|
Generation: e.Generation, Attempt: e.Attempt, Dispatched: e.Dispatched, CapacityWait: e.CapacityWait, GenerationDuration: e.GenerationDuration, RetryDelay: e.RetryDelay, Duration: e.Duration,
|
|
Outcome: string(e.Outcome), Failure: string(e.Failure), StatusCode: e.StatusCode, RetriesExhausted: e.RetriesExhausted, Attempts: e.Attempts, Retries: e.TransientRetries, Repairs: e.RepairCalls,
|
|
Usage: Usage{count(e.Usage.PromptTokens), count(e.Usage.CompletionTokens), count(e.Usage.TotalTokens), count(e.Usage.CachedTokens), count(e.Usage.CacheWriteTokens)}}
|
|
r.mu.Lock()
|
|
r.records = append(r.records, o)
|
|
r.mu.Unlock()
|
|
}
|
|
}
|
|
func (r *Recorder) snapshot() []Observation {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
return append([]Observation(nil), r.records...)
|
|
}
|
|
|
|
// totals uses only finished records. Attempt usage is retained for diagnosis,
|
|
// but is never added to the already aggregated terminal usage.
|
|
func (r *Recorder) totals() (attempts, retries, repairs uint64) {
|
|
for _, o := range r.snapshot() {
|
|
if o.Kind == string(pk.ExecutionFinished) {
|
|
attempts += o.Attempts
|
|
retries += o.Retries
|
|
repairs += o.Repairs
|
|
}
|
|
}
|
|
return
|
|
}
|