@xdarkicex/openclaw-memory-libravdb 1.3.5

package/sidecar/model/registry.go

@@ -0,0 +1,262 @@
+package model
+
+import (
+	"fmt"
+	"math"
+	"sync"
+	"time"
+
+	ort "github.com/yalue/onnxruntime_go"
+)
+
+type Task int
+
+const (
+	TaskEmbedding Task = iota
+	TaskSummarization
+)
+
+type Profile struct {
+	Name          string
+	Family        string
+	Task          Task
+	Dims          int
+	MaxCtxTokens  int
+	Quantization  string
+	Normalize     bool
+	ModelPath     string
+	TokenizerPath string
+	OrtLibPath    string
+}
+
+type MemoryPolicy struct {
+	SummarizerIdleTTL  time.Duration
+	EmbedderIdleTTL    time.Duration
+	MaxTotalModelBytes int64
+	// EvictionK calibrates model-eviction sensitivity:
+	// k = idleEligibilitySeconds * medianModelSizeBytes.
+	// Increase k to keep idle models resident longer; decrease it to evict more aggressively.
+	EvictionK float64
+}
+
+type Status struct {
+	Name             string        `json:"name"`
+	Family           string        `json:"family"`
+	Task             string        `json:"task"`
+	Loaded           bool          `json:"loaded"`
+	UseCount         int           `json:"useCount"`
+	LastAccess       time.Time     `json:"lastAccess"`
+	IdleFor          time.Duration `json:"idleFor"`
+	ReservedBytes    int64         `json:"reservedBytes"`
+	EvictionPriority float64       `json:"evictionPriority"`
+}
+
+type Registry struct {
+	mu          sync.RWMutex
+	policy      MemoryPolicy
+	runtimePath string
+	loaded      map[string]*loadedModel
+}
+
+type loadedModel struct {
+	key           string
+	profile       Profile
+	lastAccess    time.Time
+	useCount      int
+	reservedBytes int64
+	closeFn       func() error
+	encoder       *EncoderModel
+	seq2seq       *Seq2SeqModel
+}
+
+const defaultEvictionK = 1.2e11
+
+var timeNow = time.Now
+
+func DefaultMemoryPolicy() MemoryPolicy {
+	return MemoryPolicy{
+		SummarizerIdleTTL:  5 * time.Minute,
+		EmbedderIdleTTL:    30 * time.Minute,
+		MaxTotalModelBytes: 2 << 30,
+		EvictionK:          defaultEvictionK,
+	}
+}
+
+func NewRegistry(policy MemoryPolicy) *Registry {
+	return &Registry{
+		policy: policy,
+		loaded: make(map[string]*loadedModel),
+	}
+}
+
+var defaultRegistry = NewRegistry(DefaultMemoryPolicy())
+
+func DefaultRegistry() *Registry {
+	return defaultRegistry
+}
+
+func (r *Registry) Close() error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	for key, loaded := range r.loaded {
+		_ = closeLoadedModel(loaded)
+		delete(r.loaded, key)
+	}
+	if ort.IsInitialized() {
+		if err := ort.DestroyEnvironment(); err != nil {
+			return err
+		}
+	}
+	r.runtimePath = ""
+	return nil
+}
+
+func (r *Registry) Status() map[string]Status {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	now := time.Now()
+	out := make(map[string]Status, len(r.loaded))
+	for key, loaded := range r.loaded {
+		out[key] = Status{
+			Name:             loaded.profile.Name,
+			Family:           loaded.profile.Family,
+			Task:             taskName(loaded.profile.Task),
+			Loaded:           true,
+			UseCount:         loaded.useCount,
+			LastAccess:       loaded.lastAccess,
+			IdleFor:          now.Sub(loaded.lastAccess),
+			ReservedBytes:    loaded.reservedBytes,
+			EvictionPriority: evictionPriority(*loaded, now, r.policy.EvictionK),
+		}
+	}
+	return out
+}
+
+func (r *Registry) Unload(name string) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	loaded, ok := r.loaded[name]
+	if !ok {
+		return nil
+	}
+	if err := closeLoadedModel(loaded); err != nil {
+		return err
+	}
+	delete(r.loaded, name)
+	if len(r.loaded) == 0 && ort.IsInitialized() {
+		if err := ort.DestroyEnvironment(); err != nil {
+			return err
+		}
+		r.runtimePath = ""
+	}
+	return nil
+}
+
+func (r *Registry) touchLocked(key string) {
+	if loaded, ok := r.loaded[key]; ok {
+		loaded.lastAccess = time.Now()
+		loaded.useCount++
+	}
+}
+
+func (r *Registry) ensureRuntimeLocked(path string) error {
+	if path == "" {
+		return fmt.Errorf("onnx runtime path is required")
+	}
+	if ort.IsInitialized() {
+		if r.runtimePath != "" && r.runtimePath != path {
+			return fmt.Errorf("onnx runtime already initialized with %q, cannot switch to %q", r.runtimePath, path)
+		}
+		r.runtimePath = path
+		return nil
+	}
+	ort.SetSharedLibraryPath(path)
+	if err := ort.InitializeEnvironment(); err != nil {
+		return err
+	}
+	r.runtimePath = path
+	return nil
+}
+
+func (r *Registry) maybeEvictLocked(now time.Time) error {
+	for key, loaded := range r.loaded {
+		idleTTL := r.policy.EmbedderIdleTTL
+		if loaded.profile.Task == TaskSummarization {
+			idleTTL = r.policy.SummarizerIdleTTL
+		}
+		if idleTTL > 0 && now.Sub(loaded.lastAccess) > idleTTL {
+			if err := closeLoadedModel(loaded); err != nil {
+				return err
+			}
+			delete(r.loaded, key)
+		}
+	}
+
+	if r.policy.MaxTotalModelBytes <= 0 {
+		return nil
+	}
+	for totalReservedBytes(r.loaded) > r.policy.MaxTotalModelBytes {
+		evictKey := ""
+		evictScore := 0.0
+		for key, loaded := range r.loaded {
+			score := evictionPriority(*loaded, now, r.policy.EvictionK)
+			if evictKey == "" || score > evictScore {
+				evictKey = key
+				evictScore = score
+			}
+		}
+		if evictKey == "" {
+			return nil
+		}
+		if err := closeLoadedModel(r.loaded[evictKey]); err != nil {
+			return err
+		}
+		delete(r.loaded, evictKey)
+	}
+	return nil
+}
+
+func evictionPriority(m loadedModel, now time.Time, k float64) float64 {
+	deltaT := now.Sub(m.lastAccess).Seconds()
+	if deltaT < 0 {
+		deltaT = 0
+	}
+	size := float64(m.reservedBytes)
+	if size < 0 {
+		size = 0
+	}
+	if k <= 0 {
+		k = defaultEvictionK
+	}
+	freq := float64(m.useCount)
+	return (deltaT * size) / (k * (1.0 + math.Log(freq+1.0)))
+}
+
+func totalReservedBytes(loaded map[string]*loadedModel) int64 {
+	var total int64
+	for _, item := range loaded {
+		total += item.reservedBytes
+	}
+	return total
+}
+
+func closeLoadedModel(loaded *loadedModel) error {
+	if loaded == nil || loaded.closeFn == nil {
+		return nil
+	}
+	return loaded.closeFn()
+}
+
+func taskName(task Task) string {
+	switch task {
+	case TaskEmbedding:
+		return "embedding"
+	case TaskSummarization:
+		return "summarization"
+	default:
+		return "unknown"
+	}
+}

package/sidecar/model/registry_test.go

@@ -0,0 +1,102 @@
+package model
+
+import (
+	"math"
+	"testing"
+	"time"
+)
+
+func TestEvictionPriorityPrefersLargeIdleRarelyUsedModels(t *testing.T) {
+	now := time.Unix(10_000, 0)
+	k := defaultEvictionK
+
+	largeIdleRare := loadedModel{
+		lastAccess:    now.Add(-10 * time.Minute),
+		useCount:      1,
+		reservedBytes: 200 << 20,
+	}
+	smallRecentWarm := loadedModel{
+		lastAccess:    now.Add(-10 * time.Second),
+		useCount:      50_000,
+		reservedBytes: 60 << 20,
+	}
+
+	oldScore := evictionPriority(largeIdleRare, now, k)
+	recentScore := evictionPriority(smallRecentWarm, now, k)
+
+	if !(oldScore > recentScore) {
+		t.Fatalf("expected larger, older, colder model to have higher eviction score: old=%f recent=%f", oldScore, recentScore)
+	}
+}
+
+func TestEvictionPriorityDecreasesWithUseCount(t *testing.T) {
+	now := time.Unix(10_000, 0)
+	k := defaultEvictionK
+
+	cold := loadedModel{
+		lastAccess:    now.Add(-5 * time.Minute),
+		useCount:      1,
+		reservedBytes: 200 << 20,
+	}
+	warm := loadedModel{
+		lastAccess:    now.Add(-5 * time.Minute),
+		useCount:      20,
+		reservedBytes: 200 << 20,
+	}
+
+	coldScore := evictionPriority(cold, now, k)
+	warmScore := evictionPriority(warm, now, k)
+
+	if !(coldScore > warmScore) {
+		t.Fatalf("expected use-count damping to reduce eviction score: cold=%f warm=%f", coldScore, warmScore)
+	}
+}
+
+func TestEvictionPriorityHandlesZeroUseCountAndUsesDefaultK(t *testing.T) {
+	now := time.Unix(10_000, 0)
+	model := loadedModel{
+		lastAccess:    now.Add(-10 * time.Minute),
+		useCount:      0,
+		reservedBytes: 200 << 20,
+	}
+
+	score := evictionPriority(model, now, 0)
+	if score <= 0 {
+		t.Fatalf("expected positive eviction score for idle loaded model, got %f", score)
+	}
+}
+
+func TestEvictionPriorityHasLogarithmicDampingCurve(t *testing.T) {
+	now := time.Unix(10_000, 0)
+	k := defaultEvictionK
+	idle := 10 * time.Minute
+	size := int64(200 << 20)
+
+	modelForUseCount := func(useCount int) loadedModel {
+		return loadedModel{
+			lastAccess:    now.Add(-idle),
+			useCount:      useCount,
+			reservedBytes: size,
+		}
+	}
+
+	p1 := evictionPriority(modelForUseCount(1), now, k)
+	p100 := evictionPriority(modelForUseCount(100), now, k)
+	p10k := evictionPriority(modelForUseCount(10_000), now, k)
+
+	ratioLow := p1 / p100
+	ratioHigh := p100 / p10k
+
+	expectedLow := (1 + math.Log(101)) / (1 + math.Log(2))
+	expectedHigh := (1 + math.Log(10_001)) / (1 + math.Log(101))
+
+	if math.Abs(ratioLow-expectedLow) > 1e-9 {
+		t.Fatalf("unexpected low-range damping ratio: got %f want %f", ratioLow, expectedLow)
+	}
+	if math.Abs(ratioHigh-expectedHigh) > 1e-9 {
+		t.Fatalf("unexpected high-range damping ratio: got %f want %f", ratioHigh, expectedHigh)
+	}
+	if !(ratioLow > ratioHigh) {
+		t.Fatalf("expected stronger damping at low counts than high counts: low=%f high=%f", ratioLow, ratioHigh)
+	}
+}

package/sidecar/model/seq2seq.go

@@ -0,0 +1,133 @@
+package model
+
+import (
+	"fmt"
+	"strings"
+
+	ort "github.com/yalue/onnxruntime_go"
+)
+
+type Seq2SeqSpec struct {
+	Key            string
+	Profile        Profile
+	RuntimePath    string
+	ModelPath      string
+	EncoderPath    string
+	DecoderPath    string
+	TokenizerPath  string
+	EncoderInputs  []string
+	EncoderOutputs []string
+	DecoderInputs  []string
+	DecoderOutputs []string
+}
+
+type Seq2SeqModel struct {
+	key      string
+	registry *Registry
+	encoder  *ort.DynamicAdvancedSession
+	decoder  *ort.DynamicAdvancedSession
+}
+
+func (r *Registry) LoadSeq2Seq(spec Seq2SeqSpec) (*Seq2SeqModel, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if err := r.ensureRuntimeLocked(strings.TrimSpace(spec.RuntimePath)); err != nil {
+		return nil, fmt.Errorf("failed to initialize onnx runtime: %w", err)
+	}
+	if spec.Key == "" {
+		spec.Key = spec.Profile.Name
+	}
+	if loaded, ok := r.loaded[spec.Key]; ok && loaded.seq2seq != nil {
+		loaded.lastAccess = timeNow()
+		return loaded.seq2seq, nil
+	}
+
+	encoderPath := strings.TrimSpace(spec.EncoderPath)
+	decoderPath := strings.TrimSpace(spec.DecoderPath)
+	modelPath := strings.TrimSpace(spec.ModelPath)
+
+	if encoderPath == "" && decoderPath == "" && modelPath == "" {
+		return nil, fmt.Errorf("seq2seq model requires encoder, decoder, or model path")
+	}
+	if encoderPath == "" && modelPath != "" {
+		encoderPath = modelPath
+	}
+
+	var encoderSession *ort.DynamicAdvancedSession
+	var err error
+	if encoderPath != "" {
+		encoderSession, err = ort.NewDynamicAdvancedSession(encoderPath, spec.EncoderInputs, spec.EncoderOutputs, nil)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create encoder session: %w", err)
+		}
+	}
+
+	var decoderSession *ort.DynamicAdvancedSession
+	if decoderPath != "" {
+		decoderSession, err = ort.NewDynamicAdvancedSession(decoderPath, spec.DecoderInputs, spec.DecoderOutputs, nil)
+		if err != nil {
+			if encoderSession != nil {
+				_ = encoderSession.Destroy()
+			}
+			return nil, fmt.Errorf("failed to create decoder session: %w", err)
+		}
+	}
+
+	seq2seq := &Seq2SeqModel{
+		key:      spec.Key,
+		registry: r,
+		encoder:  encoderSession,
+		decoder:  decoderSession,
+	}
+	r.loaded[spec.Key] = &loadedModel{
+		key:           spec.Key,
+		profile:       spec.Profile,
+		lastAccess:    timeNow(),
+		useCount:      0,
+		reservedBytes: fileSize(encoderPath) + fileSize(decoderPath) + fileSize(spec.ModelPath) + fileSize(spec.TokenizerPath),
+		closeFn: func() error {
+			if decoderSession != nil {
+				if err := decoderSession.Destroy(); err != nil {
+					return err
+				}
+			}
+			if encoderSession != nil {
+				if err := encoderSession.Destroy(); err != nil {
+					return err
+				}
+			}
+			return nil
+		},
+		seq2seq: seq2seq,
+	}
+	if err := r.maybeEvictLocked(timeNow()); err != nil {
+		return nil, err
+	}
+	return seq2seq, nil
+}
+
+func (m *Seq2SeqModel) Touch() {
+	if m == nil || m.registry == nil {
+		return
+	}
+	m.registry.mu.Lock()
+	m.registry.touchLocked(m.key)
+	m.registry.mu.Unlock()
+}
+
+func (m *Seq2SeqModel) RunEncoder(inputs, outputs []ort.Value) error {
+	if m == nil || m.encoder == nil {
+		return fmt.Errorf("encoder session not loaded")
+	}
+	m.Touch()
+	return m.encoder.Run(inputs, outputs)
+}
+
+func (m *Seq2SeqModel) RunDecoder(inputs, outputs []ort.Value) error {
+	if m == nil || m.decoder == nil {
+		return fmt.Errorf("decoder session not loaded")
+	}
+	m.Touch()
+	return m.decoder.Run(inputs, outputs)
+}