@an5/adapters 0.2.11 → 0.2.14

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (91) hide show
  1. package/README.md +137 -5
  2. package/dist/an5Adapter.d.ts +42 -1
  3. package/dist/an5Adapter.js +146 -38
  4. package/dist/base/types.d.ts +14 -0
  5. package/dist/index.d.ts +2 -0
  6. package/dist/index.js +8 -1
  7. package/dist/sqlite/browserEngine.d.ts +13 -0
  8. package/dist/sqlite/browserEngine.js +13 -0
  9. package/dist/sqlite/engine.d.ts +5 -0
  10. package/dist/sqlite/engine.js +10 -0
  11. package/dist/sqlite/index.d.ts +1 -0
  12. package/dist/sqlite/index.js +1 -0
  13. package/dist/sqlite/vector.d.ts +155 -0
  14. package/dist/sqlite/vector.js +432 -0
  15. package/dotnet/An5Adapter.cs +99 -28
  16. package/dotnet/Base/Core.cs +14 -0
  17. package/dotnet/Mssql/MssqlEngine.cs +44 -1
  18. package/dotnet/Postgres/PostgresEngine.cs +2 -0
  19. package/dotnet/Sqlite/SqliteEngine.cs +82 -2
  20. package/dotnet/Sqlite/SqliteVectors.cs +401 -0
  21. package/golang/an5_adapter.go +84 -32
  22. package/golang/base/vectors.go +263 -0
  23. package/golang/sqlite/vector.go +248 -0
  24. package/golang/sqlite_vector_test.go +288 -0
  25. package/golang/table_client.go +67 -3
  26. package/java/pom.xml +177 -0
  27. package/java/src/main/java/an5/adapters/An5Adapter.java +374 -0
  28. package/java/src/main/java/an5/adapters/An5Aggregate.java +116 -0
  29. package/java/src/main/java/an5/adapters/An5GroupBy.java +97 -0
  30. package/java/src/main/java/an5/adapters/An5Query.java +100 -0
  31. package/java/src/main/java/an5/adapters/An5TableClient.java +1030 -0
  32. package/java/src/main/java/an5/adapters/An5ViewClient.java +107 -0
  33. package/java/src/main/java/an5/adapters/base/Dialect.java +103 -0
  34. package/java/src/main/java/an5/adapters/base/FieldMeta.java +42 -0
  35. package/java/src/main/java/an5/adapters/base/JdbcUrls.java +201 -0
  36. package/java/src/main/java/an5/adapters/base/Metadata.java +195 -0
  37. package/java/src/main/java/an5/adapters/base/RelationDef.java +26 -0
  38. package/java/src/main/java/an5/adapters/base/SqlBuilder.java +422 -0
  39. package/java/src/main/java/an5/adapters/base/SqliteVectors.java +375 -0
  40. package/java/src/main/java/an5/adapters/base/Vectors.java +123 -0
  41. package/java/src/test/java/An5AdaptersSmoke.java +733 -0
  42. package/java/src/test/java/Json.java +174 -0
  43. package/java/src/test/java/QuerySemantics.java +111 -0
  44. package/kotlin/build.gradle.kts +145 -0
  45. package/kotlin/settings.gradle.kts +10 -0
  46. package/kotlin/src/main/kotlin/an5/adapters/Aggregate.kt +89 -0
  47. package/kotlin/src/main/kotlin/an5/adapters/An5.kt +91 -0
  48. package/kotlin/src/main/kotlin/an5/adapters/Query.kt +130 -0
  49. package/kotlin/src/main/kotlin/an5/adapters/Rows.kt +166 -0
  50. package/kotlin/src/main/kotlin/an5/adapters/TableClient.kt +169 -0
  51. package/kotlin/src/main/kotlin/an5/adapters/Types.kt +77 -0
  52. package/kotlin/src/test/kotlin/an5/adapters/An5Smoke.kt +329 -0
  53. package/native/sqlite/CMakeLists.txt +12 -0
  54. package/native/sqlite/README.md +117 -0
  55. package/native/sqlite/an5_vector.c +209 -0
  56. package/package.json +33 -7
  57. package/python/an5_adapter.py +37 -6
  58. package/python/base/__init__.py +15 -0
  59. package/python/base/vectors.py +431 -0
  60. package/python/sqlite/__init__.py +2 -2
  61. package/python/sqlite/provider.py +21 -1
  62. package/python/table_client.py +129 -59
  63. package/rust/src/adapter.rs +342 -24
  64. package/rust/src/base/mod.rs +7 -0
  65. package/rust/src/base/vectors.rs +293 -0
  66. package/rust/src/lib.rs +3 -3
  67. package/scripts/sqlite-native-build.js +32 -0
  68. package/swift/Package.swift +22 -0
  69. package/swift/Sources/An5Adapters/An5Adapter.swift +160 -0
  70. package/swift/Sources/An5Adapters/Metadata.swift +183 -0
  71. package/swift/Sources/An5Adapters/Query.swift +105 -0
  72. package/swift/Sources/An5Adapters/Row+Accessors.swift +129 -0
  73. package/swift/Sources/An5Adapters/SQLBuilder.swift +359 -0
  74. package/swift/Sources/An5Adapters/SQLiteDriver.swift +301 -0
  75. package/swift/Sources/An5Adapters/SqliteVectors.swift +243 -0
  76. package/swift/Sources/An5Adapters/TableClient.swift +671 -0
  77. package/swift/Sources/An5Adapters/Types.swift +127 -0
  78. package/swift/Sources/An5Adapters/Vectors.swift +86 -0
  79. package/swift/Sources/An5Adapters/ViewClient.swift +50 -0
  80. package/swift/Sources/CSQLite/module.modulemap +8 -0
  81. package/swift/Sources/CSQLite/shim.h +9 -0
  82. package/swift/Tests/An5AdaptersTests/QuerySemanticsTests.swift +51 -0
  83. package/swift/Tests/An5AdaptersTests/RuntimeSmokeTests.swift +488 -0
  84. package/typescript/src/an5Adapter.ts +155 -31
  85. package/typescript/src/base/types.ts +15 -0
  86. package/typescript/src/index.ts +13 -0
  87. package/typescript/src/sqlite/browserEngine.ts +24 -0
  88. package/typescript/src/sqlite/engine.ts +13 -0
  89. package/typescript/src/sqlite/index.ts +1 -1
  90. package/typescript/src/sqlite/vector.ts +537 -0
  91. package/golang/sqlite/sqlite.go +0 -14
@@ -68,6 +68,9 @@ func GetRelationsForModel(modelName string) map[string]RelationDef {
68
68
  type An5Adapter struct {
69
69
  DB *sql.DB
70
70
  Dialect Dialect
71
+ // VectorSupport carries the SQLite vector settings. Set it before the first
72
+ // search to load sqlite-vec or to pin a strategy.
73
+ VectorSupport sqlite.VectorSupport
71
74
  }
72
75
 
73
76
  // NewAn5Adapter constructs an An5Adapter instance.
@@ -167,14 +170,47 @@ func (a *An5Adapter) VectorSearch(ctx context.Context, tableName string, targetV
167
170
  vecBytes, _ := json.Marshal(targetVector)
168
171
  vecStr := string(vecBytes)
169
172
 
170
- // 1. Primary path: Native database SQL vector query execution via dialect provider (mssql / postgres / sqlite)
173
+ // 1. SQLite ranks through sqlite-vec, a driver function or json_each,
174
+ // whichever this connection supports, so the table never has to leave the
175
+ // database. `database/sql` cannot register a user function, so the UDF
176
+ // strategy needs the driver to install them; see sqlite.RegisterVectorFunctions.
177
+ if a.Dialect == base.DialectSqlite {
178
+ tail := ""
179
+ if strings.TrimSpace(whereClause) != "" {
180
+ tail = " WHERE " + whereClause
181
+ }
182
+ support := a.VectorSupport
183
+ ranked, ran, err := sqlite.RunVectorSearch(
184
+ ctx,
185
+ a.DB,
186
+ a.QueryRaw,
187
+ base.QuoteTable(tableName, a.Dialect),
188
+ sqlite.NormalizeTableName(tableName),
189
+ base.QuoteIdentifier(vectorField, a.Dialect),
190
+ vectorField,
191
+ targetVector,
192
+ distanceMetric,
193
+ take,
194
+ tail,
195
+ args,
196
+ support,
197
+ )
198
+ if err != nil {
199
+ return nil, err
200
+ }
201
+ if ran {
202
+ return DecodeVectorRows(ranked, vectorField), nil
203
+ }
204
+ // No in-database strategy: score the column in Go instead.
205
+ return a.vectorSearchInMemory(ctx, tableName, targetVector, take, whereClause, vectorField, distanceMetric, args)
206
+ }
207
+
208
+ // 2. Primary path: Native database SQL vector query execution via dialect provider (mssql / postgres)
171
209
  var sqlQuery string
172
210
  if a.Dialect == base.DialectPostgres {
173
211
  sqlQuery = postgres.BuildVectorSearchQuery(tableName, vectorField, distanceMetric, dim, take, whereClause)
174
212
  } else if a.Dialect == base.DialectMssql {
175
213
  sqlQuery = mssql.BuildVectorSearchQuery(tableName, vectorField, distanceMetric, dim, take, whereClause)
176
- } else if a.Dialect == base.DialectSqlite {
177
- sqlQuery = sqlite.BuildVectorSearchQuery(tableName, vectorField, distanceMetric, dim, take, whereClause)
178
214
  }
179
215
 
180
216
  if sqlQuery != "" {
@@ -185,7 +221,7 @@ func (a *An5Adapter) VectorSearch(ctx context.Context, tableName string, targetV
185
221
  }
186
222
  }
187
223
 
188
- // 2. Secondary fallback: Fetch rows and compute vector distance in-memory if DB lacks native vector extensions
224
+ // 3. Secondary fallback: Fetch rows and compute vector distance in-memory if DB lacks native vector extensions
189
225
  fallbackQuery := "SELECT * FROM " + base.QuoteTable(tableName, a.Dialect)
190
226
  if strings.TrimSpace(whereClause) != "" {
191
227
  fallbackQuery += " WHERE " + whereClause
@@ -199,6 +235,45 @@ func (a *An5Adapter) VectorSearch(ctx context.Context, tableName string, targetV
199
235
  return VectorSearchFallback(rows, targetVector, take, vectorField, distanceMetric), nil
200
236
  }
201
237
 
238
+ // vectorSearchInMemory loads the table and scores the vector column in Go.
239
+ //
240
+ // The last resort for every provider, and the only option on a SQLite connection
241
+ // that offers neither sqlite-vec, a driver function nor JSON1. The whole table is
242
+ // loaded, so it is only reasonable for a table small enough to fit in memory.
243
+ func (a *An5Adapter) vectorSearchInMemory(ctx context.Context, tableName string, targetVector []float64, take int, whereClause string, vectorField string, distanceMetric string, args []interface{}) ([]map[string]interface{}, error) {
244
+ query := "SELECT * FROM " + base.QuoteTable(tableName, a.Dialect)
245
+ if strings.TrimSpace(whereClause) != "" {
246
+ query += " WHERE " + whereClause
247
+ }
248
+ rows, err := a.QueryRaw(ctx, query, args...)
249
+ if err != nil {
250
+ return nil, err
251
+ }
252
+ return VectorSearchFallback(rows, targetVector, take, vectorField, distanceMetric), nil
253
+ }
254
+
255
+ // DecodeVectorRows replaces every stored `VECTOR(n)` column with a decoded
256
+ // float64 slice, so a caller reads numbers rather than float32 bytes or JSON text.
257
+ func DecodeVectorRows(rows []map[string]interface{}, vectorField string) []map[string]interface{} {
258
+ for _, row := range rows {
259
+ for key, value := range row {
260
+ if key != vectorField {
261
+ continue
262
+ }
263
+ decoded := base.DecodeVector(value)
264
+ if decoded == nil {
265
+ continue
266
+ }
267
+ // A `distance` column is a real number, not a stored vector.
268
+ if key == "distance" {
269
+ continue
270
+ }
271
+ row[key] = decoded
272
+ }
273
+ }
274
+ return rows
275
+ }
276
+
202
277
  // VectorSearchFallback ranks rows in-memory.
203
278
  func VectorSearchFallback(rows []map[string]interface{}, targetVector []float64, take int, vectorField string, distanceMetric string) []map[string]interface{} {
204
279
  type scored struct {
@@ -213,36 +288,14 @@ func VectorSearchFallback(rows []map[string]interface{}, targetVector []float64,
213
288
  continue
214
289
  }
215
290
 
216
- var rowVec []float64
217
- switch v := raw.(type) {
218
- case string:
219
- _ = json.Unmarshal([]byte(v), &rowVec)
220
- case []byte:
221
- _ = json.Unmarshal(v, &rowVec)
222
- case []float64:
223
- rowVec = v
224
- case []interface{}:
225
- for _, item := range v {
226
- if f, ok := item.(float64); ok {
227
- rowVec = append(rowVec, f)
228
- }
229
- }
230
- }
231
-
232
- if len(rowVec) != len(targetVector) {
291
+ // A float32 BLOB, a legacy JSON text column and an already-decoded slice
292
+ // all read the same way.
293
+ rowVec := base.DecodeVector(raw)
294
+ dist, ok := base.VectorDistance(targetVector, rowVec, distanceMetric)
295
+ if !ok {
233
296
  continue
234
297
  }
235
298
 
236
- var dist float64
237
- switch strings.ToLower(distanceMetric) {
238
- case "euclidean":
239
- dist = base.EuclideanDistance(targetVector, rowVec)
240
- case "dot":
241
- dist = -base.DotProduct(targetVector, rowVec)
242
- default: // cosine
243
- dist = 1.0 - base.CosineSimilarity(targetVector, rowVec)
244
- }
245
-
246
299
  rCopy := make(map[string]interface{})
247
300
  for k, val := range r {
248
301
  rCopy[k] = val
@@ -311,4 +364,3 @@ func (a *An5Adapter) ExecuteProc(ctx context.Context, procName string, args ...i
311
364
  }
312
365
  return a.ExecuteRaw(ctx, sql, args...)
313
366
  }
314
-
@@ -0,0 +1,263 @@
1
+ package base
2
+
3
+ import (
4
+ "encoding/binary"
5
+ "encoding/json"
6
+ "math"
7
+ "strings"
8
+ )
9
+
10
+ // ─── SQLite vector codec ──────────────────────────────────────────────────────────
11
+ //
12
+ // SQLite has no vector type, so a `VECTOR(n)` column stores a BLOB of little-endian
13
+ // float32 values. These helpers are shared by the adapter (which reads and writes
14
+ // the column) and by the `sqlite` package (which builds the ranking SQL).
15
+ //
16
+ // BytesPerVectorFloat is the size of one stored value; the column length in bytes
17
+ // is the dimension times this, which is what the sqlite-vec strategy filters on.
18
+ const BytesPerVectorFloat = 4
19
+
20
+ // Vector field kinds the generated clients can describe.
21
+ const (
22
+ VectorStrategySqliteVec = "sqlite-vec"
23
+ VectorStrategyUdf = "udf"
24
+ VectorStrategySQL = "sql"
25
+ VectorStrategyMemory = "memory"
26
+ )
27
+
28
+ // An5VectorFunctions are the scalar functions the driver may register so SQLite
29
+ // itself ranks the column.
30
+ var An5VectorFunctions = map[string]string{
31
+ "cosine": "an5_vec_cosine",
32
+ "euclidean": "an5_vec_l2",
33
+ "dot": "an5_vec_ip",
34
+ }
35
+
36
+ // SqliteVecFunctions are the scalar functions the sqlite-vec extension provides.
37
+ var SqliteVecFunctions = map[string]string{
38
+ "cosine": "vec_distance_cosine",
39
+ "euclidean": "vec_distance_l2",
40
+ "dot": "vec_distance_ip",
41
+ }
42
+
43
+ // EncodeVector encodes a vector as the little-endian float32 BLOB the column stores.
44
+ func EncodeVector(values []float64) []byte {
45
+ bytes := make([]byte, len(values)*BytesPerVectorFloat)
46
+ for i, v := range values {
47
+ binary.LittleEndian.PutUint32(bytes[i*BytesPerVectorFloat:], math.Float32bits(float32(v)))
48
+ }
49
+ return bytes
50
+ }
51
+
52
+ // ParseVector reads a numeric sequence, or JSON text holding one.
53
+ func ParseVector(value interface{}) []float64 {
54
+ switch v := value.(type) {
55
+ case nil:
56
+ return nil
57
+ case []float64:
58
+ return v
59
+ case float32:
60
+ return []float64{float64(v)}
61
+ case float64:
62
+ return []float64{v}
63
+ case []float32:
64
+ out := make([]float64, len(v))
65
+ for i, f := range v {
66
+ out[i] = float64(f)
67
+ }
68
+ return out
69
+ case string:
70
+ trimmed := strings.TrimSpace(v)
71
+ if !strings.HasPrefix(trimmed, "[") {
72
+ return nil
73
+ }
74
+ var parsed []float64
75
+ if err := json.Unmarshal([]byte(trimmed), &parsed); err != nil {
76
+ return nil
77
+ }
78
+ return ParseVector(parsed)
79
+ case []byte:
80
+ // `database/sql` hands a BLOB column over as bytes; whether that is
81
+ // float32 storage or JSON text is decided by its length and first byte.
82
+ if len(v) > 0 && v[0] == '[' {
83
+ var parsed []float64
84
+ if err := json.Unmarshal(v, &parsed); err != nil {
85
+ return nil
86
+ }
87
+ return ParseVector(parsed)
88
+ }
89
+ return DecodeVector(v)
90
+ case []interface{}:
91
+ out := make([]float64, 0, len(v))
92
+ for _, item := range v {
93
+ f, ok := item.(float64)
94
+ if !ok {
95
+ return nil
96
+ }
97
+ out = append(out, f)
98
+ }
99
+ return out
100
+ default:
101
+ return nil
102
+ }
103
+ }
104
+
105
+ // DecodeVector decodes a stored vector: a float32 BLOB, a legacy JSON text
106
+ // column, or a sequence the driver already decoded.
107
+ func DecodeVector(value interface{}) []float64 {
108
+ if bytes, ok := value.([]byte); ok {
109
+ if len(bytes) == 0 || len(bytes)%BytesPerVectorFloat != 0 {
110
+ // Not float32 storage; it may still be JSON text in the same column.
111
+ return ParseVector(string(bytes))
112
+ }
113
+ if len(bytes) > 0 && bytes[0] == '[' {
114
+ var parsed []float64
115
+ if json.Unmarshal(bytes, &parsed) == nil { return parsed }
116
+ }
117
+ out := make([]float64, len(bytes)/BytesPerVectorFloat)
118
+ for i := range out {
119
+ out[i] = float64(math.Float32frombits(binary.LittleEndian.Uint32(bytes[i*BytesPerVectorFloat:])))
120
+ }
121
+ return out
122
+ }
123
+ return ParseVector(value)
124
+ }
125
+
126
+ // NeedsVectorEncoding reports whether a value still has to become bytes on its
127
+ // way into a `VECTOR(n)` column. `[]float64` and `[]interface{}` are the shapes
128
+ // a caller writes; text and bytes are already stored forms and pass through.
129
+ func NeedsVectorEncoding(value interface{}) bool {
130
+ switch value.(type) {
131
+ case nil, string, []byte:
132
+ return false
133
+ case []float64, []float32:
134
+ return true
135
+ case []interface{}:
136
+ return len(value.([]interface{})) > 0
137
+ default:
138
+ return false
139
+ }
140
+ }
141
+
142
+ // IsVectorField reports whether generated metadata describes a `VECTOR(n)` column.
143
+ // The metadata value is a map of field name to a small record, so both a typed
144
+ // struct and a plain map are accepted.
145
+ func IsVectorField(definition interface{}) bool {
146
+ switch def := definition.(type) {
147
+ case nil:
148
+ return false
149
+ case string:
150
+ return strings.Contains(def, "[]") || strings.HasPrefix(strings.ToUpper(strings.TrimSpace(def)), "VECTOR")
151
+ case map[string]interface{}:
152
+ if kind, _ := def["kind"].(string); kind == "vector" {
153
+ return true
154
+ }
155
+ for _, key := range []string{"sql", "type"} {
156
+ if sql, ok := def[key].(string); ok && sqlLooksVector(sql) {
157
+ return true
158
+ }
159
+ }
160
+ ts, _ := def["ts"].(string)
161
+ return strings.Contains(ts, "[]")
162
+ default:
163
+ return false
164
+ }
165
+ }
166
+
167
+ // sqlLooksVector reports whether a declared SQL type names a vector column.
168
+ func sqlLooksVector(sql string) bool {
169
+ trimmed := strings.ToUpper(strings.TrimSpace(sql))
170
+ if !strings.HasPrefix(trimmed, "VECTOR") {
171
+ return false
172
+ }
173
+ rest := trimmed[len("VECTOR"):]
174
+ return rest == "" || strings.HasPrefix(rest, "(") || strings.HasPrefix(rest, " ")
175
+ }
176
+
177
+ // PlanVectorStrategies returns the ordered strategies to try for a connection.
178
+ //
179
+ // capabilities reports what the connection supports; declaredType is the
180
+ // column's DDL type, because the JSON strategy only reaches rows stored as text
181
+ // and so a BLOB column drops it from the order. preference pins one strategy.
182
+ func PlanVectorStrategies(capabilities map[string]bool, declaredType, preference string) []string {
183
+ if preference != "" && preference != "auto" {
184
+ if preference == VectorStrategyMemory {
185
+ return nil
186
+ }
187
+ return []string{preference}
188
+ }
189
+
190
+ var available []string
191
+ if capabilities["vec"] {
192
+ available = append(available, VectorStrategySqliteVec)
193
+ }
194
+ if capabilities["udf"] {
195
+ available = append(available, VectorStrategyUdf)
196
+ }
197
+ if capabilities["json1"] {
198
+ available = append(available, VectorStrategySQL)
199
+ }
200
+
201
+ declared := strings.ToUpper(strings.TrimSpace(declaredType))
202
+ if declared != "" && !strings.Contains(declared, "TEXT") && !strings.Contains(declared, "CHAR") &&
203
+ !strings.Contains(declared, "CLOB") && !strings.Contains(declared, "STRING") &&
204
+ !strings.Contains(declared, "JSON") {
205
+ return without(available, VectorStrategySQL)
206
+ }
207
+ return available
208
+ }
209
+
210
+ func without(list []string, item string) []string {
211
+ out := make([]string, 0, len(list))
212
+ for _, entry := range list {
213
+ if entry != item {
214
+ out = append(out, entry)
215
+ }
216
+ }
217
+ return out
218
+ }
219
+
220
+ // NormalizeMetric returns a known metric name, defaulting to cosine.
221
+ func NormalizeMetric(metric string) string {
222
+ switch strings.ToLower(strings.TrimSpace(metric)) {
223
+ case "euclidean", "dot":
224
+ return strings.ToLower(strings.TrimSpace(metric))
225
+ default:
226
+ return "cosine"
227
+ }
228
+ }
229
+
230
+ // VectorDistance returns the distance between two equal-length vectors, lower is
231
+ // closer. It mirrors the SQL functions so a search ranks the same way whichever
232
+ // strategy ran. The second result is false when the pair cannot be scored, which
233
+ // is how a stored vector of another dimension is kept out of the ranking.
234
+ func VectorDistance(a, b []float64, metric string) (float64, bool) {
235
+ if len(a) == 0 || len(a) != len(b) {
236
+ return 0, false
237
+ }
238
+ switch NormalizeMetric(metric) {
239
+ case "euclidean":
240
+ return EuclideanDistance(a, b), true
241
+ case "dot":
242
+ return -DotProduct(a, b), true
243
+ default:
244
+ m1 := math.Sqrt(sumOfSquares(a))
245
+ m2 := math.Sqrt(sumOfSquares(b))
246
+ if m1 == 0 || m2 == 0 {
247
+ return 1.0, true
248
+ }
249
+ var dot float64
250
+ for i := range a {
251
+ dot += a[i] * b[i]
252
+ }
253
+ return 1.0 - dot/(m1*m2), true
254
+ }
255
+ }
256
+
257
+ func sumOfSquares(values []float64) float64 {
258
+ var sum float64
259
+ for _, v := range values {
260
+ sum += v * v
261
+ }
262
+ return sum
263
+ }
@@ -0,0 +1,248 @@
1
+ // Package sqlite provides SQLite specific dialect execution and query generation for AN5 Go Adapters.
2
+ //
3
+ // SQLite has no vector type, so a `VECTOR(n)` column stores a BLOB of little-endian
4
+ // float32 values and is ranked in one of four ways, in this order:
5
+ //
6
+ // 1. sqlite-vec native scalar distances, when the driver loaded the extension.
7
+ // 2. udf an5_vec_cosine / an5_vec_l2 / an5_vec_ip registered with the
8
+ // driver. Reads the BLOB and the legacy JSON text.
9
+ // 3. sql json_each brute force in plain SQL. Needs no user function, so it
10
+ // works with any driver, but only reaches rows stored as JSON text.
11
+ // 4. memory the column is loaded and scored in Go.
12
+ //
13
+ // Strategies 1-3 rank inside the database and transfer only the matching rows,
14
+ // which is why they are preferred over the in-memory path.
15
+ package sqlite
16
+
17
+ import (
18
+ "context"
19
+ "database/sql"
20
+ "encoding/json"
21
+ "fmt"
22
+ "strings"
23
+
24
+ "an5adapters/base"
25
+ )
26
+
27
+ // BuildVectorSearchQuery returns empty string as SQLite ranks through its own
28
+ // strategies; see RunVectorSearch for the statements it builds.
29
+ func BuildVectorSearchQuery(table, field, metric string, dim, take int, where string) string {
30
+ _ = base.QuoteTable(table, base.DialectSqlite)
31
+ _ = base.QuoteIdentifier(field, base.DialectSqlite)
32
+ // Native vector extension not present in standard sqlite; in-memory fallback is used.
33
+ return ""
34
+ }
35
+
36
+ // VectorSupport carries what a SQLite connection needs for in-database vector
37
+ // search: the path to the sqlite-vec extension to load and a strategy to pin
38
+ // instead of probing.
39
+ type VectorSupport struct {
40
+ // SqliteVecPath names the sqlite-vec extension binary. It is only used to
41
+ // report what the caller asked for; loading the extension is the driver's
42
+ // job (mattn/go-sqlite3 takes it in the DSN as `_loc`/`_x` style options, and
43
+ // modernc.org/sqlite through its own LoadExtension API).
44
+ SqliteVecPath string
45
+ // VectorStrategy pins "sqlite-vec", "udf", "sql" or "memory". Empty probes.
46
+ VectorStrategy string
47
+ }
48
+
49
+ // ScalarFuncRegistrar installs a scalar function with a driver. `database/sql`
50
+ // has no API for this, so a driver has to hand one over:
51
+ //
52
+ // sql.Register("an5-sqlite3", &sqlite3.SQLiteDriver{
53
+ // ConnectHook: func(c *sqlite3.SQLiteConn) error {
54
+ // return sqlite.RegisterVectorFunctions(c.RegisterFunc)
55
+ // },
56
+ // })
57
+ func RegisterVectorFunctions(register func(name string, impl interface{}, pure bool) error) error {
58
+ for metric, name := range base.An5VectorFunctions {
59
+ captured := metric
60
+ if err := register(name, func(a, b interface{}) interface{} {
61
+ distance, ok := base.VectorDistance(base.DecodeVector(a), base.DecodeVector(b), captured)
62
+ if !ok { return nil }
63
+ return distance
64
+ }, true); err != nil {
65
+ return err
66
+ }
67
+ }
68
+ return nil
69
+ }
70
+
71
+ // Capabilities probes what this connection can do, once per call. Both probes are
72
+ // cheap `SELECT`s; the sqlite-vec one fails on a connection without the extension,
73
+ // and the user-function one on a driver that never registered them.
74
+ func Capabilities(ctx context.Context, db *sql.DB) map[string]bool {
75
+ caps := map[string]bool{"vec": false, "udf": false, "json1": false}
76
+ if db == nil {
77
+ return caps
78
+ }
79
+
80
+ row := db.QueryRowContext(ctx, "SELECT vec_version()")
81
+ var version interface{}
82
+ if err := row.Scan(&version); err == nil {
83
+ caps["vec"] = true
84
+ }
85
+
86
+ // The probe has to reach a real distance: a driver that resolved the name
87
+ // would still fail on the arguments otherwise.
88
+ row = db.QueryRowContext(ctx, "SELECT an5_vec_cosine(zeroblob(4), zeroblob(4))")
89
+ var distance interface{}
90
+ if err := row.Scan(&distance); err == nil {
91
+ caps["udf"] = true
92
+ }
93
+
94
+ row = db.QueryRowContext(ctx, "SELECT json_valid('[1]')")
95
+ var jsonOK interface{}
96
+ if err := row.Scan(&jsonOK); err == nil {
97
+ caps["json1"] = true
98
+ }
99
+
100
+ return caps
101
+ }
102
+
103
+ // ReadColumnType returns a vector column's declared DDL type, or "" when the
104
+ // table or column is not there.
105
+ func ReadColumnType(ctx context.Context, db *sql.DB, table, column string) string {
106
+ row := db.QueryRowContext(ctx, "SELECT type FROM pragma_table_info(?) WHERE name = ?", table, column)
107
+ var declared sql.NullString
108
+ if err := row.Scan(&declared); err != nil || !declared.Valid {
109
+ return ""
110
+ }
111
+ return declared.String
112
+ }
113
+
114
+ // jsonText guards the JSON path: `json_valid` fails on a float32 BLOB.
115
+ func jsonText(column string) string {
116
+ return fmt.Sprintf("CASE WHEN json_valid(%s) THEN %s ELSE '[]' END", column, column)
117
+ }
118
+
119
+ func jsonDistanceExpr(metric, column string) string {
120
+ row := jsonText(column)
121
+ sameLength := fmt.Sprintf("(SELECT COUNT(*) FROM json_each(%s)) = (SELECT COUNT(*) FROM q)", row)
122
+
123
+ var terms string
124
+ switch metric {
125
+ case "euclidean":
126
+ terms = fmt.Sprintf("sqrt((SELECT SUM((je.value - q.v) * (je.value - q.v)) FROM json_each(%s) je JOIN q ON q.k = je.key))", row)
127
+ case "dot":
128
+ terms = fmt.Sprintf("-(SELECT SUM(je.value * q.v) FROM json_each(%s) je JOIN q ON q.k = je.key)", row)
129
+ default:
130
+ terms = fmt.Sprintf("1.0 - (SELECT SUM(je.value * q.v) FROM json_each(%s) je JOIN q ON q.k = je.key)"+
131
+ " / NULLIF(sqrt((SELECT SUM(q.v * q.v) FROM q)) * sqrt((SELECT SUM(je.value * je.value) FROM json_each(%s) je)), 0)", row, row)
132
+ }
133
+ return fmt.Sprintf("CASE WHEN %s THEN %s END", sameLength, terms)
134
+ }
135
+
136
+ // BuildRankingQuery returns the ranking SQL for one strategy, with the arguments
137
+ // the query vector needs appended to `args`.
138
+ //
139
+ // sqlite-vec and udf bind the query vector as a float32 BLOB; sql binds it as
140
+ // JSON text because json_each reads text.
141
+ //
142
+ // A row whose stored vector cannot be scored (a NULL, a text column, another
143
+ // dimension) yields a NULL distance; those rows are dropped rather than reported,
144
+ // so a caller never has to sort them out of the result.
145
+ func BuildRankingQuery(strategy, metric, table, column string, vector []float64, take int, tail string, args []interface{}) (string, []interface{}) {
146
+ rank := func(inner string) string {
147
+ return fmt.Sprintf("SELECT * FROM (%s) AS an5_ranked WHERE distance IS NOT NULL ORDER BY distance ASC LIMIT %d", inner, take)
148
+ }
149
+
150
+ switch metric {
151
+ case "euclidean", "dot":
152
+ default:
153
+ metric = "cosine"
154
+ }
155
+
156
+ if strategy == base.VectorStrategySqliteVec {
157
+ fn, ok := base.SqliteVecFunctions[metric]
158
+ if !ok {
159
+ return "", args
160
+ }
161
+ // sqlite-vec only understands float32 BLOB operands, so rows stored any
162
+ // other way are excluded instead of aborting the query. The guard joins the
163
+ // caller's WHERE with AND, since the ranking query is where both apply.
164
+ guard := fmt.Sprintf("%s IS NOT NULL", column)
165
+ // `tail` carries its own leading " WHERE", so it is appended after AND.
166
+ if tail != "" {
167
+ guard += " AND" + strings.Replace(tail, " WHERE ", " ", 1)
168
+ }
169
+ inner := fmt.Sprintf("SELECT *, CASE WHEN vec_length(vec_f32(%s)) = %d THEN %s(vec_f32(%s), ?) END AS distance FROM %s WHERE %s",
170
+ column, len(vector), fn, column, table, guard)
171
+ return rank(inner), append([]interface{}{base.EncodeVector(vector)}, args...)
172
+ }
173
+
174
+ if strategy == base.VectorStrategyUdf {
175
+ fn, ok := base.An5VectorFunctions[metric]
176
+ if !ok {
177
+ return "", args
178
+ }
179
+ inner := fmt.Sprintf("SELECT *, %s(%s, ?) AS distance FROM %s%s", fn, column, table, tail)
180
+ return rank(inner), append([]interface{}{base.EncodeVector(vector)}, args...)
181
+ }
182
+
183
+ payload, err := json.Marshal(vector)
184
+ if err != nil {
185
+ return "", args
186
+ }
187
+ inner := fmt.Sprintf("SELECT *, %s AS distance FROM %s%s", jsonDistanceExpr(metric, column), table, tail)
188
+ prefix := "WITH q AS (SELECT je.key AS k, CAST(je.value AS REAL) AS v FROM json_each(?) je) "
189
+ return prefix + rank(inner), append([]interface{}{string(payload)}, args...)
190
+ }
191
+
192
+ // RunVectorSearch ranks a SQLite table inside the database, trying each available
193
+ // strategy in order. It reports whether a strategy ran: when none does, the
194
+ // caller falls back to scoring in memory.
195
+ func RunVectorSearch(
196
+ ctx context.Context,
197
+ db *sql.DB,
198
+ exec func(ctx context.Context, query string, args ...interface{}) ([]map[string]interface{}, error),
199
+ table, rawTable, column, rawColumn string,
200
+ vector []float64,
201
+ metric string,
202
+ take int,
203
+ tail string,
204
+ args []interface{},
205
+ support VectorSupport,
206
+ ) ([]map[string]interface{}, bool, error) {
207
+ if len(vector) == 0 {
208
+ return nil, false, nil
209
+ }
210
+
211
+ caps := Capabilities(ctx, db)
212
+ if db == nil {
213
+ return nil, false, nil
214
+ }
215
+ declared := ReadColumnType(ctx, db, rawTable, rawColumn)
216
+
217
+ for _, strategy := range base.PlanVectorStrategies(caps, declared, support.VectorStrategy) {
218
+ if strategy == base.VectorStrategySQL {
219
+ clause := " WHERE "
220
+ if tail != "" { clause = tail + " AND " }
221
+ binary, err := exec(ctx, "SELECT 1 FROM " + table + clause + "typeof(" + column + ") = 'blob' LIMIT 1", args...)
222
+ if err != nil || len(binary) > 0 { continue }
223
+ }
224
+ query, queryArgs := BuildRankingQuery(strategy, metric, table, column, vector, take, tail, args)
225
+ if query == "" {
226
+ continue
227
+ }
228
+ rows, err := exec(ctx, query, queryArgs...)
229
+ if err == nil {
230
+ return rows, true, nil
231
+ }
232
+ // The connection advertised the capability but the query failed, e.g. an
233
+ // extension that did not really load. The next strategy is cheaper than
234
+ // reporting the failure, and memory is the last resort.
235
+ }
236
+ return nil, false, nil
237
+ }
238
+
239
+ // NormalizeTableName reduces a mapped table name to the bare name
240
+ // `pragma_table_info` takes as a plain string: the last segment of a qualified
241
+ // name, with the quoting a generated client may have left around it removed.
242
+ func NormalizeTableName(name string) string {
243
+ trimmed := strings.TrimSpace(name)
244
+ if index := strings.LastIndex(trimmed, "."); index >= 0 {
245
+ trimmed = trimmed[index+1:]
246
+ }
247
+ return strings.Trim(trimmed, "[]`\"")
248
+ }