memoryweave/go/internal/storage/bench_test.go

// 织忆 MemoryWeave — 存储层性能基准
package storage

import (
	"fmt"
	"strings"
	"testing"

	"github.com/xiaoxue/memoryweave/internal/models"
)

// ─── Embedder 基准 ────────────────────────────────────────

func BenchmarkEmbedder_Single(b *testing.B) {
	e := &Embedder{endpoint: "http://localhost:8000/v1/embeddings"}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		e.EncodeSingle(fmt.Sprintf("benchmark query number %d with some context", i))
	}
}

func BenchmarkEmbedder_Batch10(b *testing.B) {
	e := &Embedder{endpoint: "http://localhost:8000/v1/embeddings"}
	texts := make([]string, 10)
	for i := range texts {
		texts[i] = fmt.Sprintf("benchmark text %d for batch encoding test", i)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		e.Encode(texts)
	}
}

func BenchmarkEmbedder_Batch50(b *testing.B) {
	e := &Embedder{endpoint: "http://localhost:8000/v1/embeddings"}
	texts := make([]string, 50)
	for i := range texts {
		texts[i] = fmt.Sprintf("benchmark text %d for large batch encoding test with more context", i)
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		e.Encode(texts)
	}
}

// ─── Recall Pipeline 基准 ─────────────────────────────────

func BenchmarkRecallPipeline_10Docs(b *testing.B) {
	p := NewRecallPipeline(
		&Embedder{endpoint: "http://localhost:8000/v1/embeddings"},
		NewMemLanceClient(nil),
		NewReranker("http://localhost:8001/rerank"),
	)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		p.Recall(fmt.Sprintf("query %d about system configuration", i), "shared", 10, 0.5)
	}
}

func BenchmarkRecallPipeline_50Docs(b *testing.B) {
	p := NewRecallPipeline(
		&Embedder{endpoint: "http://localhost:8000/v1/embeddings"},
		NewMemLanceClient(nil),
		NewReranker("http://localhost:8001/rerank"),
	)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		p.Recall(fmt.Sprintf("deep query %d about project memory and system facts", i), "shared", 50, 0.5)
	}
}

// ─── 向量操作基准 ────────────────────────────────────────

func BenchmarkCosineSimilarity(b *testing.B) {
	a := make([]float32, 1024)
	bVec := make([]float32, 1024)
	for i := range a {
		a[i] = float32(i) / 1024.0
		bVec[i] = float32(1024-i) / 1024.0
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		cosineSimTest(a, bVec)
	}
}

func BenchmarkMMR_Rerank(b *testing.B) {
	results := make([]models.RecallResult, 50)
	for i := range results {
		results[i] = models.RecallResult{
			ID:      fmt.Sprintf("doc-%d", i),
			Content: fmt.Sprintf("document %d with some content for reranking", i),
			Score:   0.9 - float64(i)*0.01,
		}
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		mmrRerank(results, 10, 0.5)
	}
}

// ─── 内存分配基准 ────────────────────────────────────────

func BenchmarkLargePayload_Memory(b *testing.B) {
	content := strings.Repeat("x", 10000)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = strings.ToLower(content)
	}
}

// ─── Helper ───────────────────────────────────────────────

// cosineSimTest 向量余弦相似度
func cosineSimTest(a, b []float32) float64 {
	var dot, normA, normB float64
	for i := range a {
		dot += float64(a[i]) * float64(b[i])
		normA += float64(a[i]) * float64(a[i])
		normB += float64(b[i]) * float64(b[i])
	}
	if normA == 0 || normB == 0 {
		return 0
	}
	return dot / (float64(normA) * float64(normB))
}

// mmrRerank MMR 重排序
func mmrRerank(results []models.RecallResult, k int, lambda float64) []models.RecallResult {
	if k >= len(results) {
		return results
	}
	selected := []models.RecallResult{results[0]}
	remaining := results[1:]

	for len(selected) < k {
		bestIdx := 0
		bestScore := -1.0
		for i, r := range remaining {
			maxSim := 0.0
			for _, s := range selected {
				sim := similarity(r.Content, s.Content)
				if sim > maxSim {
					maxSim = sim
				}
			}
			mmr := lambda*r.Score - (1-lambda)*maxSim
			if mmr > bestScore {
				bestScore = mmr
				bestIdx = i
			}
		}
		selected = append(selected, remaining[bestIdx])
		remaining = append(remaining[:bestIdx], remaining[bestIdx+1:]...)
	}
	return selected
}

func similarity(a, b string) float64 {
	// Jaccard 相似度
	wordsA := make(map[string]bool)
	for _, w := range strings.Fields(a) {
		wordsA[w] = true
	}
	overlap := 0
	for _, w := range strings.Fields(b) {
		if wordsA[w] {
			overlap++
		}
	}
	if len(wordsA) == 0 {
		return 0
	}
	return float64(overlap) / float64(len(wordsA))
}