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.go

@@ -30,7 +30,10 @@ import (
 	"crypto/rand"
 	"database/sql"
 	"encoding/hex"
+	"encoding/json"
 	"fmt"
+	"math"
+	"sort"
 	"strings"
 
 	"go.starlark.net/starlark"
@@ -39,10 +42,11 @@ import (
 
 // DBModuleConfig holds configuration for a db module instance.
 type DBModuleConfig struct {
-	PackageID  string
-	CanWrite   bool // true when db.write permission is granted
-	DB         *sql.DB
-	IsPostgres bool
+	PackageID   string
+	CanWrite    bool // true when db.write permission is granted
+	DB          *sql.DB
+	IsPostgres  bool
+	HasPgvector bool // true when pgvector extension is available on Postgres
 }
 
 // allowedViews is the set of platform views extensions may query via db.view().
@@ -55,12 +59,13 @@ var allowedViews = map[string]bool{
 // prefixed with ext_{packageID}_. Write operations are guarded by CanWrite.
 func BuildDBModule(ctx context.Context, cfg DBModuleConfig) *starlarkstruct.Module {
 	fns := starlark.StringDict{
-		"query":       starlark.NewBuiltin("db.query", dbQuery(ctx, cfg)),
-		"count":       starlark.NewBuiltin("db.count", dbCount(ctx, cfg)),
-		"aggregate":   starlark.NewBuiltin("db.aggregate", dbAggregate(ctx, cfg)),
-		"query_batch": starlark.NewBuiltin("db.query_batch", dbQueryBatch(ctx, cfg)),
-		"view":        starlark.NewBuiltin("db.view", dbView(ctx, cfg)),
-		"list_tables": starlark.NewBuiltin("db.list_tables", dbListTables(ctx, cfg)),
+		"query":         starlark.NewBuiltin("db.query", dbQuery(ctx, cfg)),
+		"count":         starlark.NewBuiltin("db.count", dbCount(ctx, cfg)),
+		"aggregate":     starlark.NewBuiltin("db.aggregate", dbAggregate(ctx, cfg)),
+		"query_batch":   starlark.NewBuiltin("db.query_batch", dbQueryBatch(ctx, cfg)),
+		"query_similar": starlark.NewBuiltin("db.query_similar", dbQuerySimilar(ctx, cfg)),
+		"view":          starlark.NewBuiltin("db.view", dbView(ctx, cfg)),
+		"list_tables":   starlark.NewBuiltin("db.list_tables", dbListTables(ctx, cfg)),
 	}
 
 	if cfg.CanWrite {
@@ -160,6 +165,20 @@ func starlarkToGoValue(v starlark.Value) (any, error) {
 		return bool(val), nil
 	case starlark.NoneType:
 		return nil, nil
+	case *starlark.List:
+		goSlice := make([]any, val.Len())
+		for i := 0; i < val.Len(); i++ {
+			elem, err := starlarkToGoValue(val.Index(i))
+			if err != nil {
+				return nil, fmt.Errorf("list[%d]: %w", i, err)
+			}
+			goSlice[i] = elem
+		}
+		b, err := json.Marshal(goSlice)
+		if err != nil {
+			return nil, fmt.Errorf("list marshal: %w", err)
+		}
+		return string(b), nil
 	default:
 		return nil, fmt.Errorf("unsupported type %s", v.Type())
 	}
@@ -926,3 +945,271 @@ func dbDelete(ctx context.Context, cfg DBModuleConfig) func(*starlark.Thread, *s
 		return starlark.True, nil
 	}
 }
+
+// ── Vector Similarity ─────────────────────────
+
+// starlarkListToFloats converts a Starlark list to a Go float64 slice.
+func starlarkListToFloats(list *starlark.List) ([]float64, error) {
+	n := list.Len()
+	if n == 0 {
+		return nil, fmt.Errorf("vector must not be empty")
+	}
+	out := make([]float64, n)
+	for i := 0; i < n; i++ {
+		switch v := list.Index(i).(type) {
+		case starlark.Float:
+			out[i] = float64(v)
+		case starlark.Int:
+			i64, _ := v.Int64()
+			out[i] = float64(i64)
+		default:
+			return nil, fmt.Errorf("vector[%d]: expected number, got %s", i, list.Index(i).Type())
+		}
+	}
+	return out, nil
+}
+
+// floatsToVectorString serializes a float64 slice as a JSON array string.
+// e.g. [0.1, 0.2, 0.3] → "[0.1,0.2,0.3]"
+func floatsToVectorString(v []float64) string {
+	b, _ := json.Marshal(v)
+	return string(b)
+}
+
+// parseVectorJSON parses a JSON array string (or pgvector text) into float64 slice.
+func parseVectorJSON(s string) ([]float64, error) {
+	var out []float64
+	if err := json.Unmarshal([]byte(s), &out); err != nil {
+		return nil, fmt.Errorf("parse vector: %w", err)
+	}
+	return out, nil
+}
+
+// cosineDistance computes cosine distance between two vectors.
+// Returns 0.0 for identical vectors, 1.0 for orthogonal, 2.0 for opposite.
+func cosineDistance(a, b []float64) float64 {
+	if len(a) != len(b) || len(a) == 0 {
+		return 1.0
+	}
+	var dot, normA, normB float64
+	for i := range a {
+		dot += a[i] * b[i]
+		normA += a[i] * a[i]
+		normB += b[i] * b[i]
+	}
+	if normA == 0 || normB == 0 {
+		return 1.0
+	}
+	return 1.0 - (dot / (math.Sqrt(normA) * math.Sqrt(normB)))
+}
+
+// dbQuerySimilar implements db.query_similar(table, column, vector=[], limit=10, filters=None, metric="cosine").
+// Returns rows ordered by distance with an injected _distance float.
+//
+// Three dispatch paths:
+//   - pgvector: native <=> operator with HNSW index
+//   - fallback: fetch rows, compute cosine distance in Go, sort, return top N
+func dbQuerySimilar(ctx context.Context, cfg DBModuleConfig) func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error) {
+	return func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+		var table, column string
+		var vectorVal *starlark.List
+		var limit starlark.Int = starlark.MakeInt(10)
+		var filters starlark.Value = starlark.None
+		var metric starlark.String = "cosine"
+
+		if err := starlark.UnpackArgs(b.Name(), args, kwargs,
+			"table", &table,
+			"column", &column,
+			"vector", &vectorVal,
+			"limit?", &limit,
+			"filters?", &filters,
+			"metric?", &metric,
+		); err != nil {
+			return nil, err
+		}
+
+		if string(metric) != "cosine" {
+			return nil, fmt.Errorf("db.query_similar: unsupported metric %q (only \"cosine\" is supported)", string(metric))
+		}
+
+		queryVec, err := starlarkListToFloats(vectorVal)
+		if err != nil {
+			return nil, fmt.Errorf("db.query_similar: %w", err)
+		}
+
+		lim, ok := limit.Int64()
+		if !ok || lim < 1 {
+			lim = 10
+		}
+		if lim > 100 {
+			lim = 100
+		}
+
+		if strings.ContainsAny(column, " \t\n\"';-") {
+			return nil, fmt.Errorf("db.query_similar: invalid column name %q", column)
+		}
+
+		physTable, err := cfg.physicalTable(table)
+		if err != nil {
+			return nil, err
+		}
+
+		if cfg.HasPgvector {
+			return querySimilarPgvector(ctx, cfg, physTable, column, queryVec, int(lim), filters)
+		}
+		return querySimilarFallback(ctx, cfg, physTable, column, queryVec, int(lim), filters)
+	}
+}
+
+// querySimilarPgvector uses the native pgvector <=> operator.
+func querySimilarPgvector(ctx context.Context, cfg DBModuleConfig, physTable, column string, queryVec []float64, limit int, filters starlark.Value) (starlark.Value, error) {
+	vecStr := floatsToVectorString(queryVec)
+
+	whereClause, whereArgs, err := cfg.starlarkFiltersToSQL(filters, 1)
+	if err != nil {
+		return nil, err
+	}
+
+	nextIdx := len(whereArgs) + 1
+	vecPH := cfg.ph(nextIdx)
+	limitPH := cfg.ph(nextIdx + 1)
+
+	query := fmt.Sprintf("SELECT *, (%s <=> %s::vector) AS _distance FROM %s",
+		column, vecPH, physTable)
+	if whereClause != "" {
+		query += " " + whereClause
+	}
+	query += fmt.Sprintf(" ORDER BY %s <=> %s::vector LIMIT %s", column, vecPH, limitPH)
+
+	allArgs := append(whereArgs, vecStr, limit)
+
+	rows, err := cfg.DB.QueryContext(ctx, query, allArgs...)
+	if err != nil {
+		return nil, fmt.Errorf("db.query_similar: %w", err)
+	}
+	defer rows.Close()
+
+	return rowsToStarlark(rows)
+}
+
+// querySimilarFallback fetches rows and computes cosine distance in Go.
+// Used for SQLite (TEXT) and Postgres without pgvector (JSONB).
+func querySimilarFallback(ctx context.Context, cfg DBModuleConfig, physTable, column string, queryVec []float64, limit int, filters starlark.Value) (starlark.Value, error) {
+	whereClause, whereArgs, err := cfg.starlarkFiltersToSQL(filters, 1)
+	if err != nil {
+		return nil, err
+	}
+
+	// Fetch up to 1000 rows for distance computation.
+	fetchLimit := limit * 10
+	if fetchLimit > 1000 {
+		fetchLimit = 1000
+	}
+	if fetchLimit < limit {
+		fetchLimit = limit
+	}
+
+	limitPH := cfg.ph(len(whereArgs) + 1)
+	query := fmt.Sprintf("SELECT * FROM %s", physTable)
+	if whereClause != "" {
+		query += " " + whereClause
+	}
+	query += " LIMIT " + limitPH
+	allArgs := append(whereArgs, fetchLimit)
+
+	rows, err := cfg.DB.QueryContext(ctx, query, allArgs...)
+	if err != nil {
+		return nil, fmt.Errorf("db.query_similar: %w", err)
+	}
+	defer rows.Close()
+
+	cols, err := rows.Columns()
+	if err != nil {
+		return nil, err
+	}
+
+	// Find the vector column index.
+	vecColIdx := -1
+	for i, c := range cols {
+		if c == column {
+			vecColIdx = i
+			break
+		}
+	}
+	if vecColIdx == -1 {
+		return nil, fmt.Errorf("db.query_similar: column %q not found in table", column)
+	}
+
+	type scoredRow struct {
+		vals     []any
+		distance float64
+	}
+	var scored []scoredRow
+
+	for rows.Next() {
+		vals := make([]any, len(cols))
+		ptrs := make([]any, len(cols))
+		for i := range vals {
+			ptrs[i] = &vals[i]
+		}
+		if err := rows.Scan(ptrs...); err != nil {
+			return nil, err
+		}
+
+		// Parse the vector column value.
+		var vecStr string
+		switch v := vals[vecColIdx].(type) {
+		case string:
+			vecStr = v
+		case []byte:
+			vecStr = string(v)
+		default:
+			continue // skip rows with unparseable vectors
+		}
+
+		rowVec, err := parseVectorJSON(vecStr)
+		if err != nil {
+			continue // skip malformed vectors
+		}
+
+		dist := cosineDistance(queryVec, rowVec)
+		scored = append(scored, scoredRow{vals: vals, distance: dist})
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	// Sort by distance ascending.
+	sort.Slice(scored, func(i, j int) bool {
+		return scored[i].distance < scored[j].distance
+	})
+
+	// Take top N and build Starlark dicts.
+	if len(scored) > limit {
+		scored = scored[:limit]
+	}
+
+	var result []starlark.Value
+	for _, sr := range scored {
+		d := starlark.NewDict(len(cols) + 1)
+		for i, col := range cols {
+			sv, err := goToStarlark(sr.vals[i])
+			if err != nil {
+				return nil, fmt.Errorf("column %q: %w", col, err)
+			}
+			if err := d.SetKey(starlark.String(col), sv); err != nil {
+				return nil, err
+			}
+		}
+		// Inject _distance.
+		if err := d.SetKey(starlark.String("_distance"), starlark.Float(sr.distance)); err != nil {
+			return nil, err
+		}
+		result = append(result, d)
+	}
+
+	if result == nil {
+		result = []starlark.Value{}
+	}
+	return starlark.NewList(result), nil
+}