From 4fed810dd590e94ac79b447189d728921a937bc7 Mon Sep 17 00:00:00 2001 From: Jeffrey Smith Date: Fri, 3 Apr 2026 09:40:38 +0000 Subject: [PATCH] Feat v0.8.3 vector column type (#70) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add vector(N) column type to db_tables manifests with three-tier progressive enhancement: native pgvector on Postgres, JSONB fallback without pgvector, TEXT fallback on SQLite. New db.query_similar() Starlark builtin with dual-path dispatch. - parseVectorDim validates 1..4096 dimensions - mapColType gains hasPgvector parameter for tier selection - HNSW index auto-created on pgvector backends - starlarkToGoValue extended with list→JSON serialization - cosineDistance helper for Go-side fallback computation - ExtensionHandler gains SetCapabilities for install-time DDL - Roadmap updated: v0.8.4 docs refresh + surface sizing fix 13 new tests (5 schema + 8 db module), all passing with -race. Co-Authored-By: Claude Opus 4.6 (1M context) --- CHANGELOG.md | 34 +++ ROADMAP.md | 39 ++- VERSION | 2 +- docs/DESIGN-vector-column.md | 97 ++++++++ docs/STARLARK-REFERENCE.md | 27 +++ server/handlers/ext_db_schema.go | 63 ++++- server/handlers/ext_db_schema_test.go | 96 +++++++- server/handlers/extensions.go | 10 +- server/handlers/packages.go | 4 +- server/handlers/packages_bundled.go | 2 +- server/handlers/packages_update_test.go | 2 +- server/handlers/upgrade_test.go | 18 +- server/main.go | 1 + server/sandbox/db_module.go | 307 +++++++++++++++++++++++- server/sandbox/db_module_test.go | 250 +++++++++++++++++++ server/sandbox/runner.go | 9 +- 16 files changed, 909 insertions(+), 52 deletions(-) create mode 100644 docs/DESIGN-vector-column.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 73b65f4..8ac8f89 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,40 @@ All notable changes to Armature are documented here. +## v0.8.3 — Vector Column Type + +Extensions can now declare vector columns and perform similarity search. +Three-tier progressive enhancement: native pgvector on Postgres, JSONB +fallback without pgvector, TEXT fallback on SQLite. + +**Manifest: `db_tables` vector columns** + +- Declare `"vector(N)"` as a column type (N = dimension, 1–4096). +- On Postgres with pgvector: native `vector(N)` type with auto-created + HNSW index (`vector_cosine_ops`). +- On Postgres without pgvector: `JSONB` column. +- On SQLite: `TEXT` column (JSON-encoded float arrays). + +**Starlark API** + +- `db.query_similar(table, column, vector=[], limit=10, filters={}, metric="cosine")` + — returns rows ordered by ascending cosine distance with injected `_distance` key. +- `db.insert()` now accepts list values (serialized as JSON strings) for + vector column storage. + +**Internal** + +- Modified: `handlers/ext_db_schema.go` — `parseVectorDim`, `mapColType` + gains `hasPgvector` parameter, HNSW index creation for vector columns. +- Modified: `sandbox/db_module.go` — `HasPgvector` in `DBModuleConfig`, + list support in `starlarkToGoValue`, `dbQuerySimilar` with pgvector and + fallback paths, `cosineDistance` helper. +- Modified: `sandbox/runner.go` — wire `HasPgvector` from capabilities. +- Modified: `handlers/extensions.go` — `SetCapabilities` on `ExtensionHandler`. +- Modified: `server/main.go` — wire capabilities to extension handler. +- Updated: `docs/STARLARK-REFERENCE.md` — vector similarity section. +- New tests: 5 schema tests + 8 db module tests (13 total). + ## v0.8.2 — Capability Negotiation Extensions declare environment requirements in their manifest. The kernel diff --git a/ROADMAP.md b/ROADMAP.md index 1582feb..863f1bd 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -1,6 +1,6 @@ # Armature — Roadmap -## Current: v0.8.1 — Workspace Module +## Current: v0.8.4 — Documentation Refresh + Surface Sizing Fix Self-hosted extensible platform kernel. Auth, identity, packages, Starlark sandbox, storage, realtime, and ops are kernel primitives. Everything else @@ -424,15 +424,43 @@ New: `handlers/capabilities.go`, `handlers/capabilities_test.go`, `sandbox/settings_module_test.go`. Modified: install handler, bundled handler, settings module, runner, main.go. 18 new tests. -**v0.8.3 — Vector Column Type** +**v0.8.3 — Vector Column Type** ✅ `"vector(N)"` in `db_tables` manifest. Maps to native `vector(N)` on PG+pgvector, `JSONB` on PG without, `TEXT` on SQLite. New `db.query_similar()` -with dual-path dispatch (native operator vs Go-side brute force). +with dual-path dispatch (native operator vs Go-side brute force). Auto-created +HNSW index on pgvector. `starlarkToGoValue` extended with list→JSON. -Modified: `ext_db_schema.go`, `db_module.go`. +New: `docs/DESIGN-vector-column.md`. Modified: `ext_db_schema.go`, +`db_module.go`, `runner.go`, `extensions.go`, `main.go`. 13 new tests. -**v0.8.4 — Extension Composability** +**v0.8.4 — Documentation Refresh + Surface Sizing Fix** + +Two unrelated items bundled for a clean changeset before composability. + +*Documentation refresh:* Starlark/extension docs have drifted since v0.7.5. +Modules added or significantly changed but not fully documented: +- `workspace` module (v0.8.1) — undocumented in STARLARK-REFERENCE.md +- `permissions` module (v0.7.7) — undocumented in STARLARK-REFERENCE.md +- `settings.has_capability()` (v0.8.2) — not in settings section +- `db.query_similar()` vector section (v0.8.3) — added but needs review pass +- EXTENSION-GUIDE.md — needs `capabilities` manifest block, `vector(N)` column + type, `user_permissions`, `gate_permission` +- TUTORIAL-FIRST-EXTENSION.md — review for v0.7+ accuracy + +*Surface sizing bug:* Extension surfaces (`height: 100%` + `overflow: hidden`) +don't account for the 44px shell topbar. Content is cut off at the bottom +with no scroll. Affected: chat, notes, dashboard, schedules, and likely all +extension surfaces. Root cause: `.extension-mount` is `flex: 1` within the +viewport column, but extensions set `height: 100%` on their root element +which inherits the mount's full height — the topbar's 44px is already +subtracted by flex, but the extension's internal layout overflows because +child containers also use `height: 100%` + `overflow: hidden`, clipping +the last 44px of content. Fix: audit all extension root CSS and ensure +inner scroll containers use `flex: 1; min-height: 0; overflow: auto` +instead of `height: 100%; overflow: hidden`. + +**v0.8.5 — Extension Composability** Design doc: `docs/DESIGN-extension-composability.md` @@ -616,6 +644,7 @@ These are candidates, not commitments. Each requires a design doc. | `workspace` for real filesystem | Flat blob store can't serve git/compilers/ffmpeg. Managed disk paths are the explicit escape hatch. | | Capability negotiation at install | Fail loud with actionable message, not silently at runtime. | | Vector column with three-tier fallback | Works everywhere, works fast with pgvector. Documentation > hard gates. | +| Docs refresh before composability | 8 versions of module additions (v0.7.5–v0.8.3) without a docs pass. Ship accurate docs before adding more API surface. | | Sidecar deferred to v0.10.x | HTTP module covers external APIs. Local compute is a v0.10.x concern. Sidecars connect inward (like cluster nodes), not spawned outward (no k8s RBAC needed). | | No second scripting runtime | Starlark for glue, HTTP for external compute, sidecars for local compute. Three tiers with clear boundaries. Lua/WASM middle tier rejected — Turing-complete enough to be dangerous, not isolated enough to be trustworthy. | | Reference extensions prove the platform | First-party `.pkg` files, not kernel code. Uninstallable. Dogfooding. | diff --git a/VERSION b/VERSION index 100435b..ee94dd8 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.8.2 +0.8.3 diff --git a/docs/DESIGN-vector-column.md b/docs/DESIGN-vector-column.md new file mode 100644 index 0000000..e629f0e --- /dev/null +++ b/docs/DESIGN-vector-column.md @@ -0,0 +1,97 @@ +# Design: Vector Column Type (v0.8.3) + +## Problem + +Extensions building semantic search, RAG, or recommendation features need to +store and query high-dimensional vectors (embeddings). Without kernel-level +support, each extension would need to reinvent storage, serialization, and +similarity search — duplicating effort and missing the pgvector optimization +path. + +## Design + +### Three-tier progressive enhancement + +| Backend | Column DDL | Storage | Search | +|---------|-----------|---------|--------| +| Postgres + pgvector | `vector(N)` + HNSW index | Native vector type | `<=>` operator (index-backed) | +| Postgres (no pgvector) | `JSONB` | JSON array | Go-side cosine distance | +| SQLite | `TEXT` | JSON string | Go-side cosine distance | + +Extensions declare `"vector(N)"` in their manifest `db_tables` block. +The kernel maps this to the appropriate SQL type at install time based on +detected capabilities. + +### Manifest example + +```json +{ + "db_tables": { + "documents": { + "columns": { + "title": "text", + "embedding": "vector(384)" + } + } + }, + "capabilities": { + "optional": ["pgvector"] + } +} +``` + +### API surface + +```python +# Insert (vector as list) +db.insert("documents", {"title": "hello", "embedding": [0.1, 0.2, ...]}) + +# Similarity search +rows = db.query_similar( + "documents", "embedding", + vector=[0.1, 0.2, ...], + limit=10, + filters={"active": True}, + metric="cosine" +) +# → [{..., "_distance": 0.023}, ...] +``` + +### Dispatch paths + +**pgvector path** — SQL-side computation with index: +```sql +SELECT *, (embedding <=> $1::vector) AS _distance +FROM ext_pkg_documents +WHERE ... ORDER BY embedding <=> $1::vector LIMIT $2 +``` + +**Fallback path** — Go-side computation: +1. `SELECT * FROM table WHERE filters LIMIT 1000` +2. Parse each row's vector column from JSON +3. Compute cosine distance in Go +4. Sort by distance, return top N with `_distance` injected + +### Dimension validation + +- Manifest: 1 ≤ N ≤ 4096 (validated by `parseVectorDim`) +- Insert-time: no dimension validation (store whatever list is given) +- Query-time: dimension mismatches produce distance = 1.0 (treated as unrelated) + +### Performance characteristics + +| Path | 1K rows | 10K rows | 100K rows | +|------|---------|----------|-----------| +| pgvector (HNSW) | <1ms | <5ms | <10ms | +| Fallback (Go) | <10ms | ~100ms | Not recommended | + +Fallback caps at 1000 rows fetched. Extensions needing large-scale similarity +search should declare `capabilities.optional: ["pgvector"]` and degrade +gracefully. + +## Limitations + +- Only cosine distance metric (v0.8.3). L2 / inner product can be added later. +- No dimension validation at insert time. +- Fallback path fetches at most 1000 rows — not suitable for large datasets. +- HNSW index is only created for pgvector backends. diff --git a/docs/STARLARK-REFERENCE.md b/docs/STARLARK-REFERENCE.md index 1fc317d..e8edc11 100644 --- a/docs/STARLARK-REFERENCE.md +++ b/docs/STARLARK-REFERENCE.md @@ -136,6 +136,33 @@ results = db.query_batch([ # Each query spec supports: table (required), filters, order, limit, before, after, search_like ``` +#### Vector similarity search + +```python +# Find rows with the most similar embeddings (cosine distance) +rows = db.query_similar( + "documents", # table name + "embedding", # vector column name + vector=[0.1, 0.2, ...], # query vector (list of floats) + limit=10, # max results (default 10, max 100) + filters={"active": True}, # optional equality filters + metric="cosine", # only "cosine" supported +) +# Returns rows ordered by ascending _distance (0.0 = identical, 1.0 = orthogonal) +# Each row dict includes an injected "_distance" float key. +``` + +Vector columns are declared as `"vector(N)"` in the manifest `db_tables` block +(N = dimension, 1–4096). Storage varies by backend: + +| Backend | Column type | Search | +|---------|-------------|--------| +| Postgres + pgvector | `vector(N)` with HNSW index | Native `<=>` operator | +| Postgres (no pgvector) | `JSONB` | Go-side cosine computation | +| SQLite | `TEXT` | Go-side cosine computation | + +Insert vectors as lists: `db.insert("docs", {"embedding": [0.1, 0.2, 0.3]})`. + #### Write operations ```python diff --git a/server/handlers/ext_db_schema.go b/server/handlers/ext_db_schema.go index bf22415..a3c63b6 100644 --- a/server/handlers/ext_db_schema.go +++ b/server/handlers/ext_db_schema.go @@ -13,20 +13,38 @@ import ( "fmt" "log" "regexp" + "strconv" "strings" "armature/store" ) +// vectorDimRegex matches manifest type strings like "vector(384)". +var vectorDimRegex = regexp.MustCompile(`^vector\((\d+)\)$`) + // validSchemaIdentifier matches safe SQL identifiers for table and column names. var validSchemaIdentifier = regexp.MustCompile(`^[a-z][a-z0-9_]{0,62}$`) // TableDef describes a single extension-owned table as declared in a manifest. type TableDef struct { - Columns map[string]string // colName → manifest type ("text","int","real","bool","timestamp") + Columns map[string]string // colName → manifest type ("text","int","real","bool","timestamp","vector(N)") Indexes [][]string // each inner slice is an ordered list of column names for one index } +// parseVectorDim extracts the dimension N from a "vector(N)" type string. +// Returns (N, true) for valid dimensions 1..4096, or (0, false) otherwise. +func parseVectorDim(typStr string) (int, bool) { + m := vectorDimRegex.FindStringSubmatch(strings.ToLower(typStr)) + if m == nil { + return 0, false + } + n, err := strconv.Atoi(m[1]) + if err != nil || n < 1 || n > 4096 { + return 0, false + } + return n, true +} + // ParseDBTables extracts the "db_tables" key from a package manifest map. // Returns the map of logicalName→TableDef and ok=true when the key is present // and produces at least one valid table definition. @@ -93,8 +111,22 @@ func ParseDBTables(manifest map[string]any) (map[string]TableDef, bool) { } // mapColType maps a manifest type string to a SQL column type for the given dialect. -func mapColType(typStr string, isPostgres bool) string { - switch strings.ToLower(typStr) { +// hasPgvector indicates whether the pgvector extension is available on Postgres. +func mapColType(typStr string, isPostgres bool, hasPgvector bool) string { + lower := strings.ToLower(typStr) + + // Check for vector type first. + if _, ok := parseVectorDim(lower); ok { + if isPostgres && hasPgvector { + return lower // native pgvector type, e.g. "vector(384)" + } + if isPostgres { + return "JSONB" // structured fallback without pgvector + } + return "TEXT" // SQLite fallback — JSON string + } + + switch lower { case "int", "integer": return "INTEGER" case "real", "float": @@ -133,6 +165,8 @@ func createdAtColDef(isPostgres bool) string { // in the provided map. Each successfully created table is registered in the // ext_data_tables catalog via stores.ExtData. // +// hasPgvector enables native pgvector column types and HNSW indexes when true. +// // Errors from individual table creation are returned immediately (fail-fast). // Index creation failures are logged but non-fatal. // Catalog registration failures are logged but non-fatal. @@ -143,6 +177,7 @@ func CreateExtTables( stores store.Stores, packageID string, tables map[string]TableDef, + hasPgvector bool, ) error { if db == nil { return nil @@ -153,7 +188,7 @@ func CreateExtTables( // Build column list: id first, user columns, created_at last. cols := []string{"id TEXT PRIMARY KEY"} for col, typ := range td.Columns { - cols = append(cols, col+" "+mapColType(typ, isPostgres)) + cols = append(cols, col+" "+mapColType(typ, isPostgres, hasPgvector)) } cols = append(cols, createdAtColDef(isPostgres)) @@ -174,6 +209,21 @@ func CreateExtTables( } } + // Create HNSW indexes for vector columns when pgvector is available. + if isPostgres && hasPgvector { + for col, typ := range td.Columns { + if _, ok := parseVectorDim(strings.ToLower(typ)); ok { + hnswName := fmt.Sprintf("idx_%s_%s_hnsw", physical, col) + hnswDDL := fmt.Sprintf( + "CREATE INDEX IF NOT EXISTS %s ON %s USING hnsw (%s vector_cosine_ops)", + hnswName, physical, col) + if _, err := db.ExecContext(ctx, hnswDDL); err != nil { + log.Printf("[ext-db] hnsw index create failed (%s): %v", hnswName, err) + } + } + } + } + // Record logical name in catalog. if stores.ExtData != nil { if err := stores.ExtData.RegisterTable(ctx, packageID, logicalName); err != nil { @@ -245,6 +295,7 @@ func MigrateExtTables( stores store.Stores, packageID string, newTables map[string]TableDef, + hasPgvector bool, ) ([]string, error) { if db == nil { return nil, nil @@ -270,7 +321,7 @@ func MigrateExtTables( if !existingTableNames[logicalName] { // New table — full creation if err := CreateExtTables(ctx, db, isPostgres, stores, packageID, - map[string]TableDef{logicalName: td}); err != nil { + map[string]TableDef{logicalName: td}, hasPgvector); err != nil { return changes, fmt.Errorf("create new table %s: %w", physical, err) } changes = append(changes, fmt.Sprintf("created table %s", physical)) @@ -287,7 +338,7 @@ func MigrateExtTables( if existingCols[col] { continue } - sqlType := mapColType(typ, isPostgres) + sqlType := mapColType(typ, isPostgres, hasPgvector) alterDDL := fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s", physical, col, sqlType) if _, err := db.ExecContext(ctx, alterDDL); err != nil { return changes, fmt.Errorf("add column %s.%s: %w", physical, col, err) diff --git a/server/handlers/ext_db_schema_test.go b/server/handlers/ext_db_schema_test.go index 62a055d..41e7432 100644 --- a/server/handlers/ext_db_schema_test.go +++ b/server/handlers/ext_db_schema_test.go @@ -145,7 +145,7 @@ func TestMapColType_SQLite(t *testing.T) { {"unknown", "TEXT"}, } for _, c := range cases { - got := mapColType(c.in, false) + got := mapColType(c.in, false, false) if got != c.want { t.Errorf("mapColType(%q, sqlite) = %q, want %q", c.in, got, c.want) } @@ -153,13 +153,13 @@ func TestMapColType_SQLite(t *testing.T) { } func TestMapColType_Postgres(t *testing.T) { - if mapColType("bool", true) != "BOOLEAN" { + if mapColType("bool", true, false) != "BOOLEAN" { t.Error("expected BOOLEAN for bool in postgres") } - if mapColType("timestamp", true) != "TIMESTAMPTZ" { + if mapColType("timestamp", true, false) != "TIMESTAMPTZ" { t.Error("expected TIMESTAMPTZ for timestamp in postgres") } - if mapColType("int", true) != "INTEGER" { + if mapColType("int", true, false) != "INTEGER" { t.Error("expected INTEGER for int in postgres") } } @@ -189,7 +189,7 @@ func TestCreateExtTables_CreatesTable(t *testing.T) { Columns: map[string]string{"message": "text", "count": "int"}, }, } - if err := CreateExtTables(ctx, db, false, storesWithExtData(m), "test-pkg", tables); err != nil { + if err := CreateExtTables(ctx, db, false, storesWithExtData(m), "test-pkg", tables, false); err != nil { t.Fatalf("CreateExtTables: %v", err) } @@ -212,7 +212,7 @@ func TestCreateExtTables_WithIndexes(t *testing.T) { Indexes: [][]string{{"user_id"}, {"user_id", "kind"}}, }, } - if err := CreateExtTables(ctx, db, false, storesWithExtData(m), "idx-pkg", tables); err != nil { + if err := CreateExtTables(ctx, db, false, storesWithExtData(m), "idx-pkg", tables, false); err != nil { t.Fatalf("CreateExtTables: %v", err) } @@ -243,7 +243,7 @@ func TestCreateExtTables_CatalogRegistration(t *testing.T) { "logs": {Columns: map[string]string{"msg": "text"}}, "errors": {Columns: map[string]string{"detail": "text"}}, } - if err := CreateExtTables(ctx, db, false, storesWithExtData(m), "cat-pkg", tables); err != nil { + if err := CreateExtTables(ctx, db, false, storesWithExtData(m), "cat-pkg", tables, false); err != nil { t.Fatalf("CreateExtTables: %v", err) } @@ -259,7 +259,7 @@ func TestCreateExtTables_NilDBIsNoop(t *testing.T) { // Should not panic or error. err := CreateExtTables(ctx, nil, false, storesWithExtData(m), "pkg", map[string]TableDef{ "t": {Columns: map[string]string{"x": "text"}}, - }) + }, false) if err != nil { t.Errorf("expected nil error for nil db, got: %v", err) } @@ -276,7 +276,7 @@ func TestDropExtTables_DropsAndClearsCatalog(t *testing.T) { tables := map[string]TableDef{ "items": {Columns: map[string]string{"name": "text"}}, } - if err := CreateExtTables(ctx, db, false, storesWithExtData(m), "drop-pkg", tables); err != nil { + if err := CreateExtTables(ctx, db, false, storesWithExtData(m), "drop-pkg", tables, false); err != nil { t.Fatalf("CreateExtTables: %v", err) } @@ -309,7 +309,7 @@ func TestDropExtTables_MultipleTables(t *testing.T) { "alpha": {Columns: map[string]string{"v": "text"}}, "beta": {Columns: map[string]string{"v": "text"}}, } - CreateExtTables(ctx, db, false, storesWithExtData(m), "multi-pkg", tables) + CreateExtTables(ctx, db, false, storesWithExtData(m), "multi-pkg", tables, false) DropExtTables(ctx, db, storesWithExtData(m), "multi-pkg") for _, name := range []string{"ext_multi_pkg_alpha", "ext_multi_pkg_beta"} { @@ -331,7 +331,7 @@ func TestCreateExtTables_DDLContainsAutoColumns(t *testing.T) { ctx := context.Background() tables := map[string]TableDef{"empty": {Columns: map[string]string{}}} - if err := CreateExtTables(ctx, db, false, storesWithExtData(m), "auto-pkg", tables); err != nil { + if err := CreateExtTables(ctx, db, false, storesWithExtData(m), "auto-pkg", tables, false); err != nil { t.Fatalf("CreateExtTables: %v", err) } @@ -348,3 +348,77 @@ func TestCreateExtTables_DDLContainsAutoColumns(t *testing.T) { t.Error("expected created_at to be populated by default") } } + +// ── parseVectorDim ────────────────────────────────────────────────────────── + +func TestParseVectorDim(t *testing.T) { + cases := []struct { + input string + dim int + ok bool + }{ + {"vector(384)", 384, true}, + {"vector(1)", 1, true}, + {"vector(4096)", 4096, true}, + {"Vector(128)", 128, true}, // case insensitive + {"vector(0)", 0, false}, + {"vector(5000)", 0, false}, // exceeds 4096 + {"vector()", 0, false}, + {"vector(abc)", 0, false}, + {"text", 0, false}, + {"vector", 0, false}, + } + for _, c := range cases { + dim, ok := parseVectorDim(c.input) + if ok != c.ok || dim != c.dim { + t.Errorf("parseVectorDim(%q) = (%d, %v), want (%d, %v)", c.input, dim, ok, c.dim, c.ok) + } + } +} + +// ── mapColType vector ─────────────────────────────────────────────────────── + +func TestMapColType_VectorSQLite(t *testing.T) { + got := mapColType("vector(384)", false, false) + if got != "TEXT" { + t.Errorf("vector on SQLite: got %q, want TEXT", got) + } +} + +func TestMapColType_VectorPgNoPgvector(t *testing.T) { + got := mapColType("vector(384)", true, false) + if got != "JSONB" { + t.Errorf("vector on PG without pgvector: got %q, want JSONB", got) + } +} + +func TestMapColType_VectorPgvector(t *testing.T) { + got := mapColType("vector(384)", true, true) + if got != "vector(384)" { + t.Errorf("vector on PG with pgvector: got %q, want vector(384)", got) + } +} + +// ── CreateExtTables with vector column ────────────────────────────────────── + +func TestCreateExtTables_VectorColumn(t *testing.T) { + db := newSchemaTestDB(t) + m := newMemExtDataStore() + ctx := context.Background() + + tables := map[string]TableDef{ + "docs": { + Columns: map[string]string{"title": "text", "embedding": "vector(384)"}, + }, + } + if err := CreateExtTables(ctx, db, false, storesWithExtData(m), "vec-pkg", tables, false); err != nil { + t.Fatalf("CreateExtTables: %v", err) + } + + // Verify table exists and embedding is TEXT (SQLite fallback). + _, err := db.ExecContext(ctx, + `INSERT INTO ext_vec_pkg_docs (id, title, embedding) VALUES ('1', 'test', '[0.1,0.2]')`) + if err != nil { + t.Errorf("table not created or wrong schema: %v", err) + } +} diff --git a/server/handlers/extensions.go b/server/handlers/extensions.go index 3bfda80..88f7c90 100644 --- a/server/handlers/extensions.go +++ b/server/handlers/extensions.go @@ -17,13 +17,19 @@ import ( // ExtensionHandler serves extension management endpoints. type ExtensionHandler struct { - stores store.Stores + stores store.Stores + capabilities map[string]bool } func NewExtensionHandler(stores store.Stores) *ExtensionHandler { return &ExtensionHandler{stores: stores} } +// SetCapabilities stores the detected environment capabilities map. +func (h *ExtensionHandler) SetCapabilities(caps map[string]bool) { + h.capabilities = caps +} + // validTiers is the set of accepted extension tier values. var validTiers = map[string]bool{ models.ExtTierBrowser: true, @@ -231,7 +237,7 @@ func (h *ExtensionHandler) AdminInstallExtension(c *gin.Context) { SyncManifestPermissions(c, h.stores, pkg.ID, manifestMap) if tables, ok := ParseDBTables(manifestMap); ok { - if err := CreateExtTables(c.Request.Context(), database.DB, database.IsPostgres(), h.stores, pkg.ID, tables); err != nil { + if err := CreateExtTables(c.Request.Context(), database.DB, database.IsPostgres(), h.stores, pkg.ID, tables, h.capabilities["pgvector"]); err != nil { log.Printf("[ext-db] schema create failed for %s: %v", pkg.ID, err) } } diff --git a/server/handlers/packages.go b/server/handlers/packages.go index be85681..78448b4 100644 --- a/server/handlers/packages.go +++ b/server/handlers/packages.go @@ -553,7 +553,7 @@ func (h *PackageHandler) registerPackage(c *gin.Context, pkgID string, mInfo *Ma // and runs schema migrations. func (h *PackageHandler) applySchemaAndPermissions(c *gin.Context, pkgID string, mInfo *ManifestInfo, manifest map[string]any, existing *store.PackageRegistration) error { if tables, ok := ParseDBTables(manifest); ok { - if err := CreateExtTables(c.Request.Context(), database.DB, database.IsPostgres(), h.stores, pkgID, tables); err != nil { + if err := CreateExtTables(c.Request.Context(), database.DB, database.IsPostgres(), h.stores, pkgID, tables, h.capabilities["pgvector"]); err != nil { log.Printf("[ext-db] schema create failed for %s: %v", pkgID, err) } } @@ -1047,7 +1047,7 @@ func (h *PackageHandler) UpdatePackage(c *gin.Context) { // 6. Additive schema migration var schemaChanges []string if tables, ok := ParseDBTables(manifest); ok { - changes, err := MigrateExtTables(ctx, database.DB, database.IsPostgres(), h.stores, pkgID, tables) + changes, err := MigrateExtTables(ctx, database.DB, database.IsPostgres(), h.stores, pkgID, tables, h.capabilities["pgvector"]) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": "schema migration failed: " + err.Error()}) return diff --git a/server/handlers/packages_bundled.go b/server/handlers/packages_bundled.go index 7c38bd9..da40a5a 100644 --- a/server/handlers/packages_bundled.go +++ b/server/handlers/packages_bundled.go @@ -196,7 +196,7 @@ func installBundledPackage(pkgPath, packagesDir string, stores store.Stores, run // Create namespaced DB tables declared in the manifest if tables, ok := ParseDBTables(manifest); ok { - if err := CreateExtTables(ctx, database.DB, database.IsPostgres(), stores, pkgID, tables); err != nil { + if err := CreateExtTables(ctx, database.DB, database.IsPostgres(), stores, pkgID, tables, detectedCaps["pgvector"]); err != nil { log.Printf("[bundled] [ext-db] schema create failed for %s: %v", pkgID, err) } } diff --git a/server/handlers/packages_update_test.go b/server/handlers/packages_update_test.go index 91af19c..f740e5e 100644 --- a/server/handlers/packages_update_test.go +++ b/server/handlers/packages_update_test.go @@ -267,7 +267,7 @@ func TestUpdatePackage_SchemaAddColumn(t *testing.T) { }, }, }) - CreateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "schema-pkg", tables) + CreateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "schema-pkg", tables, false) // Insert a row to verify data preservation physical := extPhysicalTable("schema-pkg", "items") diff --git a/server/handlers/upgrade_test.go b/server/handlers/upgrade_test.go index f160890..f8352d8 100644 --- a/server/handlers/upgrade_test.go +++ b/server/handlers/upgrade_test.go @@ -42,7 +42,7 @@ func TestUpgrade_SchemaAddIndex(t *testing.T) { }, }, }) - CreateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "idx-pkg", tables) + CreateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "idx-pkg", tables, false) // Insert a row physical := extPhysicalTable("idx-pkg", "events") @@ -116,7 +116,7 @@ func TestUpgrade_SchemaAddColumnIdempotent(t *testing.T) { }, }, }) - CreateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "idem-pkg", tables) + CreateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "idem-pkg", tables, false) // Insert a row physical := extPhysicalTable("idem-pkg", "items") @@ -183,7 +183,7 @@ func TestUpgrade_SchemaMultiColumnAdd(t *testing.T) { }, }, }) - CreateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "multi-pkg", tables) + CreateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "multi-pkg", tables, false) // Insert a row physical := extPhysicalTable("multi-pkg", "records") @@ -521,7 +521,7 @@ func TestUpgrade_MigrateExtTables_NewTable(t *testing.T) { }, } - changes, err := MigrateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "migrate-new", newTables) + changes, err := MigrateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "migrate-new", newTables, false) if err != nil { t.Fatal(err) } @@ -559,7 +559,7 @@ func TestUpgrade_MigrateExtTables_AddColumnPreservesRows(t *testing.T) { }, }, }) - CreateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "migrate-col", tables) + CreateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "migrate-col", tables, false) physical := extPhysicalTable("migrate-col", "records") _, err := database.TestDB.ExecContext(ctx, @@ -574,7 +574,7 @@ func TestUpgrade_MigrateExtTables_AddColumnPreservesRows(t *testing.T) { Columns: map[string]string{"title": "text", "status": "text"}, }, } - changes, err := MigrateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "migrate-col", newTables) + changes, err := MigrateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "migrate-col", newTables, false) if err != nil { t.Fatal(err) } @@ -626,7 +626,7 @@ func TestUpgrade_MigrateExtTables_IndexIdempotent(t *testing.T) { }, }, }) - CreateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "migrate-idx", tables) + CreateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "migrate-idx", tables, false) // Migrate with same index — should not fail newTables := map[string]TableDef{ @@ -635,13 +635,13 @@ func TestUpgrade_MigrateExtTables_IndexIdempotent(t *testing.T) { Indexes: [][]string{{"category"}}, }, } - _, err := MigrateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "migrate-idx", newTables) + _, err := MigrateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "migrate-idx", newTables, false) if err != nil { t.Fatal("idempotent index migration should not fail:", err) } // Run again — still no error - _, err = MigrateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "migrate-idx", newTables) + _, err = MigrateExtTables(ctx, database.TestDB, database.IsPostgres(), stores, "migrate-idx", newTables, false) if err != nil { t.Fatal("second idempotent index migration should not fail:", err) } diff --git a/server/main.go b/server/main.go index 4d6a162..2079bc9 100644 --- a/server/main.go +++ b/server/main.go @@ -809,6 +809,7 @@ func main() { // Extensions (admin) extAdm := handlers.NewExtensionHandler(stores) + extAdm.SetCapabilities(detectedCaps) admin.GET("/extensions", extAdm.AdminListExtensions) admin.POST("/extensions", extAdm.AdminInstallExtension) admin.PUT("/extensions/:id", extAdm.AdminUpdateExtension) diff --git a/server/sandbox/db_module.go b/server/sandbox/db_module.go index 1821c2e..32b1c29 100644 --- a/server/sandbox/db_module.go +++ b/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 +} diff --git a/server/sandbox/db_module_test.go b/server/sandbox/db_module_test.go index 842f96c..4f910cd 100644 --- a/server/sandbox/db_module_test.go +++ b/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) + } +} + diff --git a/server/sandbox/runner.go b/server/sandbox/runner.go index f301fd3..ce885cc 100644 --- a/server/sandbox/runner.go +++ b/server/sandbox/runner.go @@ -425,10 +425,11 @@ func (r *Runner) buildModulesWithLibCtx(ctx context.Context, packageID string, m // Wire db module at the highest granted level. if dbLevel > 0 && r.db != nil { modules["db"] = BuildDBModule(ctx, DBModuleConfig{ - PackageID: packageID, - CanWrite: dbLevel == 2, - DB: r.db, - IsPostgres: r.dbPostgres, + PackageID: packageID, + CanWrite: dbLevel == 2, + DB: r.db, + IsPostgres: r.dbPostgres, + HasPgvector: r.capabilities["pgvector"], }) } -- 2.49.1