74 lines
2.6 KiB
Go
74 lines
2.6 KiB
Go
|
package observation
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync"
|
||
|
|
|
||
|
|
"go.maximumdirect.net/promptkit/internal/domain"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Group incrementally collects terminal facts for explicitly observed runs.
|
||
|
|
// Its zero value is ready; it must not be copied after first use.
|
||
|
|
type Group struct {
|
||
|
|
mu sync.Mutex
|
||
|
|
members map[string]bool
|
||
|
|
terminals int
|
||
|
|
attempts, retries, repairs uint64
|
||
|
|
usage domain.ObservedTokenUsage
|
||
|
|
}
|
||
|
|
|
||
|
|
// Observe enrolls a run and incorporates its first terminal event only.
|
||
|
|
func (g *Group) Observe(event domain.ExecutionEvent) {
|
||
|
|
g.mu.Lock()
|
||
|
|
defer g.mu.Unlock()
|
||
|
|
if g.members == nil {
|
||
|
|
g.members = make(map[string]bool)
|
||
|
|
}
|
||
|
|
terminal := g.members[event.RunID]
|
||
|
|
if _, exists := g.members[event.RunID]; !exists {
|
||
|
|
g.members[event.RunID] = false
|
||
|
|
}
|
||
|
|
if event.Kind != domain.ExecutionFinished || terminal {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
g.members[event.RunID] = true
|
||
|
|
g.attempts += event.Attempts
|
||
|
|
g.retries += event.TransientRetries
|
||
|
|
g.repairs += event.RepairCalls
|
||
|
|
if g.terminals == 0 {
|
||
|
|
complete := domain.ObservedTokenCount{Complete: true}
|
||
|
|
g.usage = domain.ObservedTokenUsage{PromptTokens: complete, CompletionTokens: complete, TotalTokens: complete, CachedTokens: complete, CacheWriteTokens: complete}
|
||
|
|
}
|
||
|
|
mergeGroupUsage(&g.usage, event.Usage)
|
||
|
|
g.terminals++
|
||
|
|
}
|
||
|
|
|
||
|
|
// Snapshot returns a detached consistent view of currently observed members.
|
||
|
|
func (g *Group) Snapshot() domain.ExecutionGroupSummary {
|
||
|
|
g.mu.Lock()
|
||
|
|
defer g.mu.Unlock()
|
||
|
|
result := domain.ExecutionGroupSummary{Members: make([]domain.ExecutionGroupMember, 0, len(g.members)), TerminalCoverageComplete: len(g.members) == g.terminals, Attempts: g.attempts, TransientRetries: g.retries, RepairCalls: g.repairs, Usage: g.usage}
|
||
|
|
for id, terminal := range g.members {
|
||
|
|
result.Members = append(result.Members, domain.ExecutionGroupMember{RunID: id, TerminalObserved: terminal})
|
||
|
|
}
|
||
|
|
if len(g.members) == 0 {
|
||
|
|
zero := domain.ObservedTokenCount{Known: true, Complete: true}
|
||
|
|
result.Usage = domain.ObservedTokenUsage{PromptTokens: zero, CompletionTokens: zero, TotalTokens: zero, CachedTokens: zero, CacheWriteTokens: zero}
|
||
|
|
}
|
||
|
|
return result
|
||
|
|
}
|
||
|
|
|
||
|
|
func mergeGroupUsage(total *domain.ObservedTokenUsage, next domain.ObservedTokenUsage) {
|
||
|
|
merge := func(a *domain.ObservedTokenCount, b domain.ObservedTokenCount) {
|
||
|
|
if b.Known {
|
||
|
|
a.Value += b.Value
|
||
|
|
}
|
||
|
|
a.Known = a.Known || b.Known
|
||
|
|
a.Complete = a.Complete && b.Complete
|
||
|
|
}
|
||
|
|
merge(&total.PromptTokens, next.PromptTokens)
|
||
|
|
merge(&total.CompletionTokens, next.CompletionTokens)
|
||
|
|
merge(&total.TotalTokens, next.TotalTokens)
|
||
|
|
merge(&total.CachedTokens, next.CachedTokens)
|
||
|
|
merge(&total.CacheWriteTokens, next.CacheWriteTokens)
|
||
|
|
}
|