@xdarkicex/openclaw-memory-libravdb 1.3.5

package/sidecar/embed/onnx_local.go

@@ -0,0 +1,319 @@
+package embed
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/sugarme/tokenizer"
+	"github.com/xDarkicex/openclaw-memory-libravdb/sidecar/model"
+)
+
+const defaultManifestName = "embedding.json"
+
+type Profile struct {
+	Backend     string `json:"backend"`
+	Family      string `json:"family,omitempty"`
+	Dimensions  int    `json:"dimensions"`
+	Normalize   bool   `json:"normalize"`
+	MaxContextTokens int `json:"maxContextTokens,omitempty"`
+	ModelPath   string `json:"modelPath,omitempty"`
+	Tokenizer   string `json:"tokenizerPath,omitempty"`
+	Fingerprint string `json:"fingerprint"`
+}
+
+type embeddingManifest struct {
+	Backend          string   `json:"backend,omitempty"`
+	Profile          string   `json:"profile,omitempty"`
+	Family           string   `json:"family,omitempty"`
+	Model            string   `json:"model"`
+	Tokenizer        string   `json:"tokenizer,omitempty"`
+	Dimensions       int      `json:"dimensions"`
+	Normalize        *bool    `json:"normalize,omitempty"`
+	InputNames       []string `json:"inputNames,omitempty"`
+	OutputName       string   `json:"outputName,omitempty"`
+	AddSpecialTokens *bool    `json:"addSpecialTokens,omitempty"`
+	Pooling          string   `json:"pooling,omitempty"`
+}
+
+type onnxLocalSpec struct {
+	Backend          string
+	Family           string
+	RuntimePath      string
+	ModelPath        string
+	TokenizerPath    string
+	Dimensions       int
+	Normalize        bool
+	InputNames       []string
+	OutputName       string
+	AddSpecialTokens bool
+	Pooling          string
+	Profile          Profile
+}
+
+type onnxLocalModel struct {
+	encoder *model.EncoderModel
+}
+
+var newONNXLocalBackend = func(spec onnxLocalSpec) (embeddingBackend, error) {
+	return newFileONNXBackend(spec)
+}
+
+func resolveONNXLocalSpec(cfg Config) (onnxLocalSpec, error) {
+	manifestPath, err := resolveManifestPath(cfg.ModelPath)
+	if err != nil {
+		return onnxLocalSpec{}, err
+	}
+
+	manifest, err := readManifest(manifestPath)
+	if err != nil {
+		return onnxLocalSpec{}, err
+	}
+	selectedProfile, hasProfile := lookupProfile(firstNonEmpty(cfg.Profile, manifest.Profile))
+
+	baseDir := filepath.Dir(manifestPath)
+	modelPath := resolveManifestAsset(baseDir, cfg.modelAssetOverride(), manifest.Model)
+	if modelPath == "" {
+		return onnxLocalSpec{}, fmt.Errorf("onnx-local manifest missing model path")
+	}
+
+	tokenizerPath := resolveManifestAsset(baseDir, cfg.tokenizerAssetOverride(), manifest.Tokenizer)
+	if tokenizerPath == "" {
+		return onnxLocalSpec{}, fmt.Errorf("onnx-local manifest missing tokenizer path")
+	}
+
+	if cfg.Dimensions > 0 && manifest.Dimensions > 0 && cfg.Dimensions != manifest.Dimensions {
+		return onnxLocalSpec{}, fmt.Errorf("onnx-local manifest dimensions %d do not match configured dimensions %d", manifest.Dimensions, cfg.Dimensions)
+	}
+
+	dimensions := manifest.Dimensions
+	if dimensions <= 0 && hasProfile {
+		dimensions = selectedProfile.Dimensions
+	}
+	if dimensions <= 0 {
+		dimensions = cfg.Dimensions
+	}
+	if dimensions <= 0 {
+		return onnxLocalSpec{}, fmt.Errorf("onnx-local dimensions must be configured in embedding.json or plugin config")
+	}
+
+	normalize := cfg.Normalize
+	if manifest.Normalize != nil {
+		normalize = *manifest.Normalize
+	} else if hasProfile {
+		normalize = selectedProfile.Normalize
+	}
+
+	inputNames := manifest.InputNames
+	if len(inputNames) == 0 {
+		inputNames = []string{"input_ids", "attention_mask", "token_type_ids"}
+	}
+
+	outputName := strings.TrimSpace(manifest.OutputName)
+	if outputName == "" {
+		outputName = "sentence_embedding"
+	}
+
+	addSpecialTokens := true
+	if manifest.AddSpecialTokens != nil {
+		addSpecialTokens = *manifest.AddSpecialTokens
+	}
+	pooling := strings.TrimSpace(manifest.Pooling)
+	if pooling == "" && outputName == "last_hidden_state" {
+		pooling = "mean"
+	}
+
+	profile := buildProfile(Profile{
+		Backend:          "onnx-local",
+		Family:           firstNonEmpty(strings.TrimSpace(manifest.Family), selectedProfile.Family, strings.TrimSpace(manifest.Backend), "onnx-local"),
+		Dimensions:       dimensions,
+		Normalize:        normalize,
+		MaxContextTokens: selectedProfile.MaxContextTokens,
+		ModelPath:        modelPath,
+		Tokenizer:        tokenizerPath,
+	})
+
+	return onnxLocalSpec{
+		Backend:          "onnx-local",
+		Family:           profile.Family,
+		RuntimePath:      strings.TrimSpace(cfg.RuntimePath),
+		ModelPath:        modelPath,
+		TokenizerPath:    tokenizerPath,
+		Dimensions:       dimensions,
+		Normalize:        normalize,
+		InputNames:       append([]string(nil), inputNames...),
+		OutputName:       outputName,
+		AddSpecialTokens: addSpecialTokens,
+		Pooling:          pooling,
+		Profile:          profile,
+	}, nil
+}
+
+func newFileONNXBackend(spec onnxLocalSpec) (embeddingBackend, error) {
+	encoder, err := model.DefaultRegistry().LoadEncoder(model.EncoderSpec{
+		Key: spec.Profile.Fingerprint,
+		Profile: model.Profile{
+			Name:          spec.Profile.Fingerprint,
+			Family:        spec.Profile.Family,
+			Task:          model.TaskEmbedding,
+			Dims:          spec.Dimensions,
+			Normalize:     spec.Normalize,
+			ModelPath:     spec.ModelPath,
+			TokenizerPath: spec.TokenizerPath,
+			OrtLibPath:    spec.RuntimePath,
+		},
+		RuntimePath:      spec.RuntimePath,
+		ModelPath:        spec.ModelPath,
+		TokenizerPath:    spec.TokenizerPath,
+		InputNames:       spec.InputNames,
+		OutputName:       spec.OutputName,
+		Dimensions:       spec.Dimensions,
+		AddSpecialTokens: spec.AddSpecialTokens,
+		Pooling:          spec.Pooling,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return miniLMBackend{
+		model: onnxLocalModel{
+			encoder: encoder,
+		},
+		normalize: spec.Normalize,
+	}, nil
+}
+
+func (m onnxLocalModel) Compute(sentence string, _ bool) ([]float32, error) {
+	if m.encoder == nil {
+		return nil, fmt.Errorf("encoder model not loaded")
+	}
+	return m.encoder.EmbedText(sentence)
+}
+
+func (m onnxLocalModel) TokenCount(sentence string, _ bool) (int, error) {
+	if m.encoder == nil {
+		return 0, fmt.Errorf("encoder model not loaded")
+	}
+	return m.encoder.TokenCount(sentence)
+}
+
+func (m onnxLocalModel) Encode(sentence string, _ bool) (*tokenizer.Encoding, error) {
+	if m.encoder == nil {
+		return nil, fmt.Errorf("encoder model not loaded")
+	}
+	return m.encoder.EncodeText(sentence, true)
+}
+
+func (m onnxLocalModel) ComputeEncoding(encoding tokenizer.Encoding) ([]float32, error) {
+	if m.encoder == nil {
+		return nil, fmt.Errorf("encoder model not loaded")
+	}
+	return m.encoder.EmbedEncoding(encoding)
+}
+
+func meanPoolLastHiddenState(flat []float32, attentionMask []int, seqLength int, dimensions int) []float32 {
+	vec := make([]float32, dimensions)
+	var denom float32
+	for tokenIdx := 0; tokenIdx < seqLength; tokenIdx++ {
+		if tokenIdx >= len(attentionMask) || attentionMask[tokenIdx] == 0 {
+			continue
+		}
+		denom++
+		base := tokenIdx * dimensions
+		for dim := 0; dim < dimensions; dim++ {
+			vec[dim] += flat[base+dim]
+		}
+	}
+	if denom == 0 {
+		return vec
+	}
+	for dim := range vec {
+		vec[dim] /= denom
+	}
+	return vec
+}
+
+func resolveManifestPath(modelPath string) (string, error) {
+	raw := strings.TrimSpace(modelPath)
+	if raw == "" {
+		return "", fmt.Errorf("onnx-local requires embeddingModelPath pointing to a model directory or embedding.json")
+	}
+
+	info, err := os.Stat(raw)
+	if err == nil && info.IsDir() {
+		return filepath.Join(raw, defaultManifestName), nil
+	}
+	if err == nil && strings.EqualFold(filepath.Base(raw), defaultManifestName) {
+		return raw, nil
+	}
+	if err == nil {
+		return filepath.Join(filepath.Dir(raw), defaultManifestName), nil
+	}
+	return "", err
+}
+
+func readManifest(path string) (embeddingManifest, error) {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return embeddingManifest{}, fmt.Errorf("failed to read %s: %w", path, err)
+	}
+	var manifest embeddingManifest
+	if err := json.Unmarshal(data, &manifest); err != nil {
+		return embeddingManifest{}, fmt.Errorf("failed to parse %s: %w", path, err)
+	}
+	return manifest, nil
+}
+
+func resolveManifestAsset(baseDir string, values ...string) string {
+	asset := firstNonEmpty(values...)
+	if asset == "" {
+		return ""
+	}
+	if filepath.IsAbs(asset) {
+		return asset
+	}
+	return filepath.Join(baseDir, asset)
+}
+
+func buildProfile(profile Profile) Profile {
+	hash := sha256.Sum256([]byte(strings.Join([]string{
+		profile.Backend,
+		profile.Family,
+		fmt.Sprintf("%d", profile.Dimensions),
+		fmt.Sprintf("%t", profile.Normalize),
+		fmt.Sprintf("%d", profile.MaxContextTokens),
+		profile.ModelPath,
+		profile.Tokenizer,
+	}, "\x00")))
+	profile.Fingerprint = hex.EncodeToString(hash[:])
+	return profile
+}
+
+func firstNonEmpty(values ...string) string {
+	for _, value := range values {
+		value = strings.TrimSpace(value)
+		if value != "" {
+			return value
+		}
+	}
+	return ""
+}
+
+func (cfg Config) modelAssetOverride() string {
+	modelPath := strings.TrimSpace(cfg.ModelPath)
+	if strings.EqualFold(filepath.Base(modelPath), defaultManifestName) {
+		return ""
+	}
+	if strings.HasSuffix(strings.ToLower(modelPath), ".onnx") {
+		return modelPath
+	}
+	return ""
+}
+
+func (cfg Config) tokenizerAssetOverride() string {
+	return strings.TrimSpace(cfg.TokenizerPath)
+}

package/sidecar/embed/onnx_local_test.go

@@ -0,0 +1,159 @@
+package embed
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+)
+
+type fakeONNXLocalBackend struct {
+	vector []float32
+}
+
+func (b fakeONNXLocalBackend) Embed(_ string, dimensions int) ([]float32, error) {
+	if len(b.vector) != dimensions {
+		return nil, nil
+	}
+	return append([]float32(nil), b.vector...), nil
+}
+
+func TestONNXLocalLoadsManifestDirectory(t *testing.T) {
+	dir := t.TempDir()
+	err := os.WriteFile(filepath.Join(dir, "embedding.json"), []byte(`{
+		"backend":"onnx-local",
+		"profile":"nomic-embed-text-v1.5",
+		"family":"nomic-embed-text-v1.5",
+		"model":"model.onnx",
+		"tokenizer":"tokenizer.json",
+		"dimensions":768,
+		"normalize":true,
+		"inputNames":["input_ids","attention_mask"],
+		"outputName":"sentence_embedding"
+	}`), 0o644)
+	if err != nil {
+		t.Fatalf("write manifest: %v", err)
+	}
+
+	original := newONNXLocalBackend
+	t.Cleanup(func() { newONNXLocalBackend = original })
+	newONNXLocalBackend = func(spec onnxLocalSpec) (embeddingBackend, error) {
+		if spec.Family != "nomic-embed-text-v1.5" {
+			t.Fatalf("unexpected family %q", spec.Family)
+		}
+		if !strings.HasSuffix(spec.ModelPath, "model.onnx") {
+			t.Fatalf("unexpected model path %q", spec.ModelPath)
+		}
+		if !strings.HasSuffix(spec.TokenizerPath, "tokenizer.json") {
+			t.Fatalf("unexpected tokenizer path %q", spec.TokenizerPath)
+		}
+		if spec.Dimensions != 768 {
+			t.Fatalf("unexpected dimensions %d", spec.Dimensions)
+		}
+		if len(spec.InputNames) != 2 {
+			t.Fatalf("unexpected input names %#v", spec.InputNames)
+		}
+		if spec.RuntimePath != "/opt/onnx/libonnxruntime.so" {
+			t.Fatalf("unexpected runtime path %q", spec.RuntimePath)
+		}
+		return fakeONNXLocalBackend{vector: make([]float32, 768)}, nil
+	}
+
+	engine := NewWithConfig(Config{
+		Backend:     "onnx-local",
+		RuntimePath: "/opt/onnx/libonnxruntime.so",
+		ModelPath:   dir,
+		Normalize:   true,
+	})
+	if !engine.Ready() {
+		t.Fatalf("expected onnx-local engine to be ready, reason=%q", engine.Reason())
+	}
+	if engine.Mode() != "onnx-local" {
+		t.Fatalf("unexpected mode %q", engine.Mode())
+	}
+	if engine.Dimensions() != 768 {
+		t.Fatalf("expected manifest dimensions 768, got %d", engine.Dimensions())
+	}
+	if engine.Profile().Fingerprint == "" {
+		t.Fatalf("expected embedding fingerprint")
+	}
+}
+
+func TestONNXLocalRejectsManifestDimensionMismatch(t *testing.T) {
+	dir := t.TempDir()
+	err := os.WriteFile(filepath.Join(dir, "embedding.json"), []byte(`{
+		"model":"model.onnx",
+		"tokenizer":"tokenizer.json",
+		"dimensions":768
+	}`), 0o644)
+	if err != nil {
+		t.Fatalf("write manifest: %v", err)
+	}
+
+	original := newONNXLocalBackend
+	t.Cleanup(func() { newONNXLocalBackend = original })
+	newONNXLocalBackend = func(spec onnxLocalSpec) (embeddingBackend, error) {
+		return fakeONNXLocalBackend{vector: make([]float32, spec.Dimensions)}, nil
+	}
+
+	engine := NewWithConfig(Config{
+		Backend:     "onnx-local",
+		RuntimePath: "/opt/onnx/libonnxruntime.so",
+		ModelPath:   dir,
+		Dimensions:  384,
+	})
+	if engine.Ready() {
+		t.Fatalf("expected mismatched dimensions to fail readiness")
+	}
+}
+
+func TestONNXLocalCanUseShippedProfileDefaults(t *testing.T) {
+	dir := t.TempDir()
+	err := os.WriteFile(filepath.Join(dir, "embedding.json"), []byte(`{
+		"profile":"nomic-embed-text-v1.5",
+		"model":"model.onnx",
+		"tokenizer":"tokenizer.json"
+	}`), 0o644)
+	if err != nil {
+		t.Fatalf("write manifest: %v", err)
+	}
+
+	original := newONNXLocalBackend
+	t.Cleanup(func() { newONNXLocalBackend = original })
+	newONNXLocalBackend = func(spec onnxLocalSpec) (embeddingBackend, error) {
+		if spec.Dimensions != 768 {
+			t.Fatalf("expected profile dimensions 768, got %d", spec.Dimensions)
+		}
+		if spec.Family != "nomic-embed-text-v1.5" {
+			t.Fatalf("expected profile family, got %q", spec.Family)
+		}
+		return fakeONNXLocalBackend{vector: make([]float32, 768)}, nil
+	}
+
+	engine := NewWithConfig(Config{
+		Backend:     "onnx-local",
+		RuntimePath: "/opt/onnx/libonnxruntime.so",
+		ModelPath:   dir,
+	})
+	if !engine.Ready() {
+		t.Fatalf("expected profile-backed onnx-local engine to be ready, reason=%q", engine.Reason())
+	}
+}
+
+func TestMeanPoolLastHiddenStateUsesAttentionMask(t *testing.T) {
+	flat := []float32{
+		1, 2,
+		3, 4,
+		100, 200,
+	}
+	mask := []int{1, 1, 0}
+
+	got := meanPoolLastHiddenState(flat, mask, 3, 2)
+
+	if len(got) != 2 {
+		t.Fatalf("unexpected pooled vector length %d", len(got))
+	}
+	if got[0] != 2 || got[1] != 3 {
+		t.Fatalf("unexpected pooled vector %#v", got)
+	}
+}

package/sidecar/embed/profile_contract_test.go

@@ -0,0 +1,71 @@
+package embed
+
+import (
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+type tokenizerContract struct {
+	Truncation *struct {
+		MaxLength int `json:"max_length"`
+	} `json:"truncation"`
+}
+
+func TestProfileMaxCtxMatchesTokenizerContract(t *testing.T) {
+	cases := []struct {
+		profile         string
+		modelDir        string
+		expectEnforced  bool
+		expectedMax     int
+	}{
+		{
+			profile:        "all-minilm-l6-v2",
+			modelDir:       filepath.Clean(filepath.Join("..", "..", ".models", "all-minilm-l6-v2")),
+			expectEnforced: true,
+			expectedMax:    128,
+		},
+		{
+			profile:        "nomic-embed-text-v1.5",
+			modelDir:       filepath.Clean(filepath.Join("..", "..", ".models", "nomic-embed-text-v1.5")),
+			expectEnforced: false,
+			expectedMax:    8192,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.profile, func(t *testing.T) {
+			profile, ok := lookupProfile(tc.profile)
+			if !ok {
+				t.Fatalf("profile %s not found", tc.profile)
+			}
+			if profile.MaxContextTokens != tc.expectedMax {
+				t.Fatalf("profile %s MaxContextTokens=%d want %d", tc.profile, profile.MaxContextTokens, tc.expectedMax)
+			}
+
+			data, err := os.ReadFile(filepath.Join(tc.modelDir, "tokenizer.json"))
+			if err != nil {
+				t.Fatalf("read tokenizer.json: %v", err)
+			}
+			var contract tokenizerContract
+			if err := json.Unmarshal(data, &contract); err != nil {
+				t.Fatalf("parse tokenizer.json: %v", err)
+			}
+
+			if tc.expectEnforced {
+				if contract.Truncation == nil {
+					t.Fatalf("expected tokenizer truncation contract for %s", tc.profile)
+				}
+				if contract.Truncation.MaxLength != tc.expectedMax {
+					t.Fatalf("profile %s MaxContextTokens=%d but tokenizer enforces %d", tc.profile, tc.expectedMax, contract.Truncation.MaxLength)
+				}
+				return
+			}
+
+			if contract.Truncation != nil {
+				t.Fatalf("expected tokenizer truncation to be disabled for %s, got max_length=%d", tc.profile, contract.Truncation.MaxLength)
+			}
+		})
+	}
+}