Feat v0.8.3 vector column (#70)

Co-authored-by: Jeffrey Smith <jasafpro@gmail.com>
Co-committed-by: Jeffrey Smith <jasafpro@gmail.com>

server/sandbox/db_module_test.go

@@ -3,6 +3,7 @@ package sandbox
 import (
 	"context"
 	"database/sql"
+	"math"
 	"strings"
 	"testing"
 
@@ -684,3 +685,252 @@ func TestDBQueryBatch_MissingTable(t *testing.T) {
 	}
 }
 
+// ── starlarkToGoValue list ──────────────────
+
+func TestStarlarkToGoValue_List(t *testing.T) {
+	list := starlark.NewList([]starlark.Value{
+		starlark.Float(0.1), starlark.Float(0.2), starlark.Float(0.3),
+	})
+	got, err := starlarkToGoValue(list)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	s, ok := got.(string)
+	if !ok {
+		t.Fatalf("expected string, got %T", got)
+	}
+	if s != "[0.1,0.2,0.3]" {
+		t.Errorf("got %q, want %q", s, "[0.1,0.2,0.3]")
+	}
+}
+
+// ── cosineDistance ───────────────────────────
+
+func TestCosineDistance(t *testing.T) {
+	// Identical vectors → distance 0
+	d := cosineDistance([]float64{1, 0, 0}, []float64{1, 0, 0})
+	if math.Abs(d) > 1e-9 {
+		t.Errorf("identical vectors: got %f, want 0", d)
+	}
+
+	// Orthogonal vectors → distance 1
+	d = cosineDistance([]float64{1, 0}, []float64{0, 1})
+	if math.Abs(d-1.0) > 1e-9 {
+		t.Errorf("orthogonal vectors: got %f, want 1", d)
+	}
+
+	// Opposite vectors → distance 2
+	d = cosineDistance([]float64{1, 0}, []float64{-1, 0})
+	if math.Abs(d-2.0) > 1e-9 {
+		t.Errorf("opposite vectors: got %f, want 2", d)
+	}
+
+	// Empty/mismatched → distance 1
+	d = cosineDistance([]float64{}, []float64{})
+	if d != 1.0 {
+		t.Errorf("empty vectors: got %f, want 1", d)
+	}
+	d = cosineDistance([]float64{1}, []float64{1, 2})
+	if d != 1.0 {
+		t.Errorf("mismatched vectors: got %f, want 1", d)
+	}
+}
+
+// ── vector helper: newVectorTestDB ──────────
+
+func newVectorTestDB(t *testing.T) (*sql.DB, DBModuleConfig) {
+	t.Helper()
+	db, err := sql.Open("sqlite", ":memory:")
+	if err != nil {
+		t.Fatalf("open db: %v", err)
+	}
+	t.Cleanup(func() { db.Close() })
+
+	// Create a table with a TEXT column for vector storage (SQLite fallback).
+	_, err = db.Exec(`CREATE TABLE ext_test_ext_embeddings (
+		id         TEXT PRIMARY KEY,
+		embedding  TEXT,
+		category   TEXT,
+		created_at TEXT DEFAULT (datetime('now'))
+	)`)
+	if err != nil {
+		t.Fatalf("create test table: %v", err)
+	}
+
+	cfg := DBModuleConfig{
+		PackageID:   "test-ext",
+		CanWrite:    true,
+		DB:          db,
+		IsPostgres:  false,
+		HasPgvector: false,
+	}
+	return db, cfg
+}
+
+// ── db.query_similar ────────────────────────
+
+func TestDBQuerySimilar_Basic(t *testing.T) {
+	db, cfg := newVectorTestDB(t)
+
+	// Insert 5 rows with known vectors.
+	db.Exec(`INSERT INTO ext_test_ext_embeddings (id, embedding, category) VALUES ('a', '[1,0,0]', 'x')`)
+	db.Exec(`INSERT INTO ext_test_ext_embeddings (id, embedding, category) VALUES ('b', '[0,1,0]', 'x')`)
+	db.Exec(`INSERT INTO ext_test_ext_embeddings (id, embedding, category) VALUES ('c', '[0,0,1]', 'x')`)
+	db.Exec(`INSERT INTO ext_test_ext_embeddings (id, embedding, category) VALUES ('d', '[0.9,0.1,0]', 'x')`)
+	db.Exec(`INSERT INTO ext_test_ext_embeddings (id, embedding, category) VALUES ('e', '[0,0.9,0.1]', 'x')`)
+
+	// Query similar to [1,0,0] — should return 'a' first (identical), then 'd' (close).
+	result, err := execScript(t, cfg, `
+rows = db.query_similar("embeddings", "embedding", vector=[1.0, 0.0, 0.0], limit=3)
+`)
+	if err != nil {
+		t.Fatalf("script error: %v", err)
+	}
+
+	rows, ok := result.Globals["rows"].(*starlark.List)
+	if !ok {
+		t.Fatalf("rows is %T, want *starlark.List", result.Globals["rows"])
+	}
+	if rows.Len() != 3 {
+		t.Fatalf("got %d rows, want 3", rows.Len())
+	}
+
+	// First row should be 'a' (distance ≈ 0).
+	first := rows.Index(0).(*starlark.Dict)
+	idVal, _, _ := first.Get(starlark.String("id"))
+	if string(idVal.(starlark.String)) != "a" {
+		t.Errorf("first row id = %s, want 'a'", idVal)
+	}
+
+	// Verify _distance is present and near 0.
+	distVal, _, _ := first.Get(starlark.String("_distance"))
+	dist := float64(distVal.(starlark.Float))
+	if dist > 0.01 {
+		t.Errorf("first row _distance = %f, want ≈ 0", dist)
+	}
+
+	// Second row should be 'd' (closest after identical).
+	second := rows.Index(1).(*starlark.Dict)
+	idVal2, _, _ := second.Get(starlark.String("id"))
+	if string(idVal2.(starlark.String)) != "d" {
+		t.Errorf("second row id = %s, want 'd'", idVal2)
+	}
+}
+
+func TestDBQuerySimilar_WithFilters(t *testing.T) {
+	db, cfg := newVectorTestDB(t)
+
+	db.Exec(`INSERT INTO ext_test_ext_embeddings (id, embedding, category) VALUES ('a', '[1,0,0]', 'alpha')`)
+	db.Exec(`INSERT INTO ext_test_ext_embeddings (id, embedding, category) VALUES ('b', '[0.9,0.1,0]', 'beta')`)
+	db.Exec(`INSERT INTO ext_test_ext_embeddings (id, embedding, category) VALUES ('c', '[0,1,0]', 'alpha')`)
+
+	// Filter to alpha only — should exclude 'b' even though it's close.
+	result, err := execScript(t, cfg, `
+rows = db.query_similar("embeddings", "embedding", vector=[1.0, 0.0, 0.0], filters={"category": "alpha"})
+`)
+	if err != nil {
+		t.Fatalf("script error: %v", err)
+	}
+
+	rows := result.Globals["rows"].(*starlark.List)
+	for i := 0; i < rows.Len(); i++ {
+		d := rows.Index(i).(*starlark.Dict)
+		idVal, _, _ := d.Get(starlark.String("id"))
+		if string(idVal.(starlark.String)) == "b" {
+			t.Error("filtered query should not include row 'b' (category=beta)")
+		}
+	}
+}
+
+func TestDBQuerySimilar_EmptyTable(t *testing.T) {
+	_, cfg := newVectorTestDB(t)
+
+	result, err := execScript(t, cfg, `
+rows = db.query_similar("embeddings", "embedding", vector=[1.0, 0.0, 0.0])
+`)
+	if err != nil {
+		t.Fatalf("script error: %v", err)
+	}
+	rows := result.Globals["rows"].(*starlark.List)
+	if rows.Len() != 0 {
+		t.Errorf("expected 0 rows for empty table, got %d", rows.Len())
+	}
+}
+
+func TestDBQuerySimilar_InvalidMetric(t *testing.T) {
+	_, cfg := newVectorTestDB(t)
+
+	_, err := execScript(t, cfg, `
+rows = db.query_similar("embeddings", "embedding", vector=[1.0], metric="l2")
+`)
+	if err == nil {
+		t.Fatal("expected error for unsupported metric")
+	}
+	if !strings.Contains(err.Error(), "unsupported metric") {
+		t.Errorf("unexpected error: %v", err)
+	}
+}
+
+func TestDBQuerySimilar_LimitCap(t *testing.T) {
+	db, cfg := newVectorTestDB(t)
+
+	// Insert 5 rows
+	for i := 0; i < 5; i++ {
+		db.Exec(`INSERT INTO ext_test_ext_embeddings (id, embedding) VALUES (?, '[1,0,0]')`,
+			generateID())
+	}
+
+	result, err := execScript(t, cfg, `
+rows = db.query_similar("embeddings", "embedding", vector=[1.0, 0.0, 0.0], limit=2)
+`)
+	if err != nil {
+		t.Fatalf("script error: %v", err)
+	}
+	rows := result.Globals["rows"].(*starlark.List)
+	if rows.Len() != 2 {
+		t.Errorf("expected 2 rows with limit=2, got %d", rows.Len())
+	}
+}
+
+func TestDBInsert_VectorColumn(t *testing.T) {
+	db, cfg := newVectorTestDB(t)
+
+	_, err := execScript(t, cfg, `
+row = db.insert("embeddings", {"embedding": [0.1, 0.2, 0.3], "category": "test"})
+`)
+	if err != nil {
+		t.Fatalf("script error: %v", err)
+	}
+
+	// Verify the stored value is valid JSON.
+	var stored string
+	db.QueryRow(`SELECT embedding FROM ext_test_ext_embeddings WHERE category = 'test'`).Scan(&stored)
+	if stored != "[0.1,0.2,0.3]" {
+		t.Errorf("stored vector = %q, want %q", stored, "[0.1,0.2,0.3]")
+	}
+}
+
+func TestDBQuerySimilar_InsertThenQuery(t *testing.T) {
+	_, cfg := newVectorTestDB(t)
+
+	// Full round-trip: insert via db.insert, then query_similar.
+	result, err := execScript(t, cfg, `
+db.insert("embeddings", {"embedding": [1.0, 0.0, 0.0], "category": "a"})
+db.insert("embeddings", {"embedding": [0.0, 1.0, 0.0], "category": "b"})
+rows = db.query_similar("embeddings", "embedding", vector=[1.0, 0.0, 0.0], limit=1)
+`)
+	if err != nil {
+		t.Fatalf("script error: %v", err)
+	}
+
+	rows := result.Globals["rows"].(*starlark.List)
+	if rows.Len() != 1 {
+		t.Fatalf("expected 1 row, got %d", rows.Len())
+	}
+	first := rows.Index(0).(*starlark.Dict)
+	catVal, _, _ := first.Get(starlark.String("category"))
+	if string(catVal.(starlark.String)) != "a" {
+		t.Errorf("closest match should be category=a, got %s", catVal)
+	}
+}
+