feat: QualityTracker.Snapshot/Restore + Router.QualityTracker() for cross-session persistence
This commit is contained in:
70
internal/router/quality_json.go
Normal file
70
internal/router/quality_json.go
Normal file
@@ -0,0 +1,70 @@
|
||||
package router
|
||||
|
||||
// QualitySnapshot is the serializable form of QualityTracker EMA data.
|
||||
type QualitySnapshot struct {
|
||||
Scores map[string]map[string]*EMAScore `json:"scores"` // ArmID -> TaskType.String() -> score
|
||||
}
|
||||
|
||||
// Snapshot returns a serializable copy of current quality data.
|
||||
func (qt *QualityTracker) Snapshot() QualitySnapshot {
|
||||
qt.mu.RLock()
|
||||
defer qt.mu.RUnlock()
|
||||
snap := QualitySnapshot{
|
||||
Scores: make(map[string]map[string]*EMAScore),
|
||||
}
|
||||
for armID, tasks := range qt.scores {
|
||||
snap.Scores[string(armID)] = make(map[string]*EMAScore)
|
||||
for taskType, score := range tasks {
|
||||
snap.Scores[string(armID)][taskType.String()] = &EMAScore{
|
||||
Value: score.Value,
|
||||
Count: score.Count,
|
||||
}
|
||||
}
|
||||
}
|
||||
return snap
|
||||
}
|
||||
|
||||
// Restore loads previously persisted quality data.
|
||||
// Replaces all current scores.
|
||||
func (qt *QualityTracker) Restore(snap QualitySnapshot) {
|
||||
qt.mu.Lock()
|
||||
defer qt.mu.Unlock()
|
||||
qt.scores = make(map[ArmID]map[TaskType]*EMAScore)
|
||||
for armID, tasks := range snap.Scores {
|
||||
qt.scores[ArmID(armID)] = make(map[TaskType]*EMAScore)
|
||||
for taskTypeStr, score := range tasks {
|
||||
taskType, ok := parseTaskType(taskTypeStr)
|
||||
if !ok {
|
||||
continue // skip unrecognized task types
|
||||
}
|
||||
qt.scores[ArmID(armID)][taskType] = &EMAScore{
|
||||
Value: score.Value,
|
||||
Count: score.Count,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// parseTaskType maps a TaskType string representation back to its constant.
|
||||
func parseTaskType(s string) (TaskType, bool) {
|
||||
for _, tt := range allTaskTypes {
|
||||
if tt.String() == s {
|
||||
return tt, true
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
// allTaskTypes lists every known TaskType constant for reverse lookup.
|
||||
var allTaskTypes = []TaskType{
|
||||
TaskBoilerplate,
|
||||
TaskGeneration,
|
||||
TaskRefactor,
|
||||
TaskReview,
|
||||
TaskUnitTest,
|
||||
TaskPlanning,
|
||||
TaskOrchestration,
|
||||
TaskSecurityReview,
|
||||
TaskDebug,
|
||||
TaskExplain,
|
||||
}
|
||||
74
internal/router/quality_json_test.go
Normal file
74
internal/router/quality_json_test.go
Normal file
@@ -0,0 +1,74 @@
|
||||
package router_test
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"somegit.dev/Owlibou/gnoma/internal/router"
|
||||
)
|
||||
|
||||
func TestQualityTracker_SnapshotRestore_RoundTrip(t *testing.T) {
|
||||
qt := router.NewQualityTracker()
|
||||
// Record some outcomes
|
||||
qt.Record("anthropic/claude-3-5-sonnet", router.TaskGeneration, true)
|
||||
qt.Record("anthropic/claude-3-5-sonnet", router.TaskGeneration, true)
|
||||
qt.Record("anthropic/claude-3-5-sonnet", router.TaskGeneration, false)
|
||||
qt.Record("ollama/gemma3", router.TaskBoilerplate, true)
|
||||
|
||||
snap := qt.Snapshot()
|
||||
|
||||
// Verify snapshot has the data
|
||||
if len(snap.Scores) == 0 {
|
||||
t.Fatal("snapshot scores should not be empty")
|
||||
}
|
||||
|
||||
// Marshal and unmarshal to simulate disk persistence
|
||||
data, err := json.Marshal(snap)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var restored router.QualitySnapshot
|
||||
if err := json.Unmarshal(data, &restored); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Restore into a fresh tracker
|
||||
qt2 := router.NewQualityTracker()
|
||||
qt2.Restore(restored)
|
||||
|
||||
// After restore, Quality() should return data (Count >= minObservations=3)
|
||||
score, hasData := qt2.Quality("anthropic/claude-3-5-sonnet", router.TaskGeneration)
|
||||
if !hasData {
|
||||
t.Error("expected quality data after restore")
|
||||
}
|
||||
if score <= 0 {
|
||||
t.Errorf("expected positive score, got %f", score)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQualityTracker_Snapshot_Empty(t *testing.T) {
|
||||
qt := router.NewQualityTracker()
|
||||
snap := qt.Snapshot()
|
||||
if snap.Scores == nil {
|
||||
t.Error("scores map should be initialized (not nil)")
|
||||
}
|
||||
if len(snap.Scores) != 0 {
|
||||
t.Errorf("expected empty scores, got %d", len(snap.Scores))
|
||||
}
|
||||
}
|
||||
|
||||
func TestQualityTracker_Restore_Replaces(t *testing.T) {
|
||||
qt := router.NewQualityTracker()
|
||||
qt.Record("arm-a", router.TaskDebug, true)
|
||||
qt.Record("arm-a", router.TaskDebug, true)
|
||||
qt.Record("arm-a", router.TaskDebug, true)
|
||||
|
||||
// Restore with different data — old data should be gone
|
||||
empty := router.QualitySnapshot{Scores: make(map[string]map[string]*router.EMAScore)}
|
||||
qt.Restore(empty)
|
||||
|
||||
_, hasData := qt.Quality("arm-a", router.TaskDebug)
|
||||
if hasData {
|
||||
t.Error("old data should be gone after restore with empty snapshot")
|
||||
}
|
||||
}
|
||||
@@ -172,6 +172,11 @@ func (r *Router) LookupArm(id ArmID) (*Arm, bool) {
|
||||
return arm, ok
|
||||
}
|
||||
|
||||
// QualityTracker returns the router's quality tracker for persistence.
|
||||
func (r *Router) QualityTracker() *QualityTracker {
|
||||
return r.quality
|
||||
}
|
||||
|
||||
// Arms returns all registered arms.
|
||||
func (r *Router) Arms() []*Arm {
|
||||
r.mu.RLock()
|
||||
|
||||
Reference in New Issue
Block a user