feat: solve issues #9 and #10 (typo tolerance + neighbor-aware embeddings)

docs/reference/api.md

@@ -23,6 +23,7 @@ func NewCombinedMatcher(embedder Embedder) ElementMatcher
 // Standalone matchers
 func NewLexicalMatcher() ElementMatcher
 func NewEmbeddingMatcher(e Embedder) ElementMatcher
+func NewEmbeddingMatcherWithNeighborWeight(e Embedder, weight float64) ElementMatcher
 
 // Built-in embedder (feature hashing, zero deps)
 func NewHashingEmbedder(dim int) Embedder

internal/engine/benchmark_test.go

@@ -76,6 +76,24 @@ func BenchmarkLexicalFind_200Elements(b *testing.B) {
 	}
 }
 
+func BenchmarkLexicalFind_TypoQuery_200Elements(b *testing.B) {
+	m := NewLexicalMatcher()
+	base := benchElements()
+	elements := make([]types.ElementDescriptor, 0, 200)
+	for len(elements) < 200 {
+		elements = append(elements, base...)
+	}
+	elements = elements[:200]
+
+	ctx := context.Background()
+	opts := types.FindOptions{Threshold: 0.0, TopK: 3}
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, _ = m.Find(ctx, "setttings", elements, opts)
+	}
+}
+
 func BenchmarkHashingEmbed(b *testing.B) {
 	h := NewHashingEmbedder(128)
 	texts := []string{"sign in button"}
@@ -108,6 +126,18 @@ func BenchmarkEmbeddingFind(b *testing.B) {
 	}
 }
 
+func BenchmarkEmbeddingFind_NoNeighborContext(b *testing.B) {
+	m := NewEmbeddingMatcherWithNeighborWeight(NewHashingEmbedder(128), 0)
+	elements := benchElements()
+	ctx := context.Background()
+	opts := types.FindOptions{Threshold: 0.3, TopK: 3}
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, _ = m.Find(ctx, "sign in button", elements, opts)
+	}
+}
+
 func BenchmarkCombinedFind(b *testing.B) {
 	m := NewCombinedMatcher(NewHashingEmbedder(128))
 	elements := benchElements()

internal/engine/embedding.go

@@ -20,12 +20,27 @@ type Embedder interface {
 // EmbeddingMatcher scores elements using cosine similarity on dense
 // vectors produced by an Embedder.
 type EmbeddingMatcher struct {
-	embedder Embedder
+	embedder       Embedder
+	neighborWeight float32
 }
 
+const defaultNeighborWeight float32 = 0.10
+
 // NewEmbeddingMatcher creates an embedding-based matcher.
 func NewEmbeddingMatcher(e Embedder) *EmbeddingMatcher {
-	return &EmbeddingMatcher{embedder: e}
+	return NewEmbeddingMatcherWithNeighborWeight(e, float64(defaultNeighborWeight))
+}
+
+// NewEmbeddingMatcherWithNeighborWeight creates an embedding matcher and sets
+// how much immediate neighbors influence each element embedding.
+func NewEmbeddingMatcherWithNeighborWeight(e Embedder, weight float64) *EmbeddingMatcher {
+	if weight < 0 {
+		weight = 0
+	}
+	if weight > 1 {
+		weight = 1
+	}
+	return &EmbeddingMatcher{embedder: e, neighborWeight: float32(weight)}
 }
 
 func (m *EmbeddingMatcher) Strategy() string {
@@ -52,6 +67,10 @@ func (m *EmbeddingMatcher) Find(_ context.Context, query string, elements []type
 
 	queryVec := vectors[0]
 	elemVecs := vectors[1:]
+	contextVecs := elemVecs
+	if m.neighborWeight > 0 && len(elemVecs) > 1 {
+		contextVecs = m.withNeighborContext(elemVecs)
+	}
 
 	type scored struct {
 		desc  types.ElementDescriptor
@@ -60,7 +79,7 @@ func (m *EmbeddingMatcher) Find(_ context.Context, query string, elements []type
 
 	var candidates []scored
 	for i, el := range elements {
-		sim := CosineSimilarity(queryVec, elemVecs[i])
+		sim := CosineSimilarity(queryVec, contextVecs[i])
 		if sim >= opts.Threshold {
 			candidates = append(candidates, scored{desc: el, score: sim})
 		}
@@ -96,6 +115,43 @@ func (m *EmbeddingMatcher) Find(_ context.Context, query string, elements []type
 	return result, nil
 }
 
+func (m *EmbeddingMatcher) withNeighborContext(base [][]float32) [][]float32 {
+	contextual := make([][]float32, len(base))
+	for i := range base {
+		vec := make([]float32, len(base[i]))
+		copy(vec, base[i])
+
+		if i > 0 {
+			for d := range vec {
+				vec[d] += base[i-1][d] * m.neighborWeight
+			}
+		}
+		if i+1 < len(base) {
+			for d := range vec {
+				vec[d] += base[i+1][d] * m.neighborWeight
+			}
+		}
+
+		normalizeDenseVector(vec)
+		contextual[i] = vec
+	}
+	return contextual
+}
+
+func normalizeDenseVector(vec []float32) {
+	var norm float64
+	for _, v := range vec {
+		norm += float64(v) * float64(v)
+	}
+	if norm == 0 {
+		return
+	}
+	invNorm := float32(1.0 / math.Sqrt(norm))
+	for i := range vec {
+		vec[i] *= invNorm
+	}
+}
+
 // CosineSimilarity computes the cosine similarity between two float32 vectors.
 func CosineSimilarity(a, b []float32) float64 {
 	if len(a) != len(b) || len(a) == 0 {

internal/engine/embedding_test.go

@@ -2,6 +2,7 @@ package engine
 
 import (
 	"context"
+	"fmt"
 	"github.com/pinchtab/semantic/internal/types"
 	"math"
 	"testing"
@@ -112,6 +113,25 @@ func TestEmbeddingMatcher_Strategy(t *testing.T) {
 	}
 }
 
+func TestNewEmbeddingMatcher_DefaultNeighborWeight(t *testing.T) {
+	m := NewEmbeddingMatcher(newDummyEmbedder(64))
+	if math.Abs(float64(m.neighborWeight-defaultNeighborWeight)) > 1e-6 {
+		t.Errorf("expected default neighborWeight=%f, got %f", defaultNeighborWeight, m.neighborWeight)
+	}
+}
+
+func TestNewEmbeddingMatcherWithNeighborWeight_ClampsRange(t *testing.T) {
+	below := NewEmbeddingMatcherWithNeighborWeight(newDummyEmbedder(64), -0.25)
+	if below.neighborWeight != 0 {
+		t.Errorf("expected neighborWeight to clamp to 0, got %f", below.neighborWeight)
+	}
+
+	above := NewEmbeddingMatcherWithNeighborWeight(newDummyEmbedder(64), 2)
+	if above.neighborWeight != 1 {
+		t.Errorf("expected neighborWeight to clamp to 1, got %f", above.neighborWeight)
+	}
+}
+
 func TestEmbeddingMatcher_Find(t *testing.T) {
 	m := NewEmbeddingMatcher(newDummyEmbedder(64))
 
@@ -167,4 +187,111 @@ func TestEmbeddingMatcher_ThresholdFiltering(t *testing.T) {
 	}
 }
 
+func TestEmbeddingMatcher_NeighborContextDisambiguatesRealWorldButtons(t *testing.T) {
+	e := newScriptedEmbedder(map[string][]float32{
+		"laptop add to cart":         {1, 1, 0},
+		"heading: Gaming Laptop Pro": {0, 1, 0},
+		"button: Add to cart":        {1, 0, 0},
+		"heading: Running Shoes":     {0, 0, 1},
+	})
+
+	elements := []types.ElementDescriptor{
+		{Ref: "title-laptop", Role: "heading", Name: "Gaming Laptop Pro"},
+		{Ref: "add-laptop", Role: "button", Name: "Add to cart"},
+		{Ref: "title-shoes", Role: "heading", Name: "Running Shoes"},
+		{Ref: "add-shoes", Role: "button", Name: "Add to cart"},
+	}
+
+	noCtx := NewEmbeddingMatcherWithNeighborWeight(e, 0)
+	baseRes, err := noCtx.Find(context.Background(), "laptop add to cart", elements, types.FindOptions{Threshold: 0, TopK: 4})
+	if err != nil {
+		t.Fatalf("Find with no context failed: %v", err)
+	}
+	baseLaptop, ok := findMatchScore(baseRes.Matches, "add-laptop")
+	if !ok {
+		t.Fatalf("expected add-laptop result in no-context run")
+	}
+	baseShoes, ok := findMatchScore(baseRes.Matches, "add-shoes")
+	if !ok {
+		t.Fatalf("expected add-shoes result in no-context run")
+	}
+	if math.Abs(baseLaptop-baseShoes) > 1e-9 {
+		t.Fatalf("expected identical button scores without context, got laptop=%.6f shoes=%.6f", baseLaptop, baseShoes)
+	}
+
+	withCtx := NewEmbeddingMatcherWithNeighborWeight(e, 0.2)
+	ctxRes, err := withCtx.Find(context.Background(), "laptop add to cart", elements, types.FindOptions{Threshold: 0, TopK: 4})
+	if err != nil {
+		t.Fatalf("Find with context failed: %v", err)
+	}
+	ctxLaptop, ok := findMatchScore(ctxRes.Matches, "add-laptop")
+	if !ok {
+		t.Fatalf("expected add-laptop result in contextual run")
+	}
+	ctxShoes, ok := findMatchScore(ctxRes.Matches, "add-shoes")
+	if !ok {
+		t.Fatalf("expected add-shoes result in contextual run")
+	}
+	if ctxLaptop <= ctxShoes {
+		t.Fatalf("expected laptop button to rank higher with context, got laptop=%.6f shoes=%.6f", ctxLaptop, ctxShoes)
+	}
+}
+
+func TestEmbeddingMatcher_SingleElement_WithNeighborWeight(t *testing.T) {
+	e := newScriptedEmbedder(map[string][]float32{
+		"open account settings":    {1, 1, 0},
+		"button: Account settings": {1, 1, 0},
+	})
+	m := NewEmbeddingMatcherWithNeighborWeight(e, 0.5)
+
+	res, err := m.Find(context.Background(), "open account settings", []types.ElementDescriptor{
+		{Ref: "settings", Role: "button", Name: "Account settings"},
+	}, types.FindOptions{Threshold: 0, TopK: 1})
+	if err != nil {
+		t.Fatalf("Find failed: %v", err)
+	}
+	if res.BestRef != "settings" {
+		t.Fatalf("expected BestRef=settings, got %s", res.BestRef)
+	}
+	if len(res.Matches) != 1 {
+		t.Fatalf("expected one match, got %d", len(res.Matches))
+	}
+}
+
+type scriptedEmbedder struct {
+	vectors map[string][]float32
+}
+
+func newScriptedEmbedder(vectors map[string][]float32) *scriptedEmbedder {
+	return &scriptedEmbedder{vectors: vectors}
+}
+
+func (s *scriptedEmbedder) Strategy() string {
+	return "scripted"
+}
+
+func (s *scriptedEmbedder) Embed(texts []string) ([][]float32, error) {
+	out := make([][]float32, len(texts))
+	for i, text := range texts {
+		base, ok := s.vectors[text]
+		if !ok {
+			return nil, fmt.Errorf("missing scripted embedding for %q", text)
+		}
+		vec := make([]float32, len(base))
+		copy(vec, base)
+		normalizeDenseVector(vec)
+		out[i] = vec
+	}
+	return out, nil
+}
+
+func findMatchScore(matches []types.ElementMatch, ref string) (float64, bool) {
+	for _, match := range matches {
+		if match.Ref == ref {
+			return match.Score, true
+		}
+	}
+	return 0, false
+}
+
 // FindResult.ConfidenceLabel tests