1
// Copyright (c) 2023 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
 
5
// @ts-expect-error - no type declarations for node:assert in this context
6
import * as assert from 'node:assert';
7
import { KnownModel, DistanceMetric } from 'cloudflare:vectorize';
8
 
9
/**
10
 * @typedef {{'vector-search': Vectorize}} Env
11
 *
12
 */
13
 
14
export const test_vector_search_vector_query = {
15
  /**
16
   * @param {unknown} _
17
   * @param {Env} env
18
   */
19
  async test(_, env) {
20
    const IDX = env['vector-search'];
21
    {
22
      // with returnValues = true, returnMetadata = "indexed"
23
      const results = await IDX.query(new Float32Array(5), {
24
        topK: 3,
25
        returnValues: true,
26
        returnMetadata: 'indexed',
27
      });
28
      assert.equal(true, results.count > 0);
29
      /** @type {VectorizeMatches}  */
30
      const expected = {
31
        matches: [
32
          {
33
            id: 'b0daca4a-ffd8-4865-926b-e24800af2a2d',
34
            values: [0.2331, 1.0125, 0.6131, 0.9421, 0.9661, 0.8121],
35
            metadata: { text: 'She sells seashells by the seashore' },
36
            score: 0.71151,
37
          },
38
          {
39
            id: 'a44706aa-a366-48bc-8cc1-3feffd87d548',
40
            values: [0.2321, 0.8121, 0.6315, 0.6151, 0.4121, 0.1512],
41
            metadata: {
42
              text: 'Peter Piper picked a peck of pickled peppers',
43
            },
44
            score: 0.68913,
45
          },
46
          {
47
            id: '43cfcb31-07e2-411f-8bf9-f82a95ba8b96',
48
            values: [0.0515, 0.7512, 0.8612, 0.2153, 0.15121, 0.6812],
49
            metadata: {
50
              text: 'You know New York, you need New York, you know you need unique New York',
51
            },
52
            score: 0.94812,
53
          },
54
        ],
55
        count: 3,
56
      };
57
      assert.deepStrictEqual(results, expected);
58
    }
59
 
60
    {
61
      // with returnValues = true, returnMetadata = "indexed"
62
      const results = await IDX.queryById('some-vector-id', {
63
        topK: 3,
64
        returnValues: true,
65
        returnMetadata: 'indexed',
66
      });
67
      assert.equal(true, results.count > 0);
68
      /** @type {VectorizeMatches}  */
69
      const expected = {
70
        matches: [
71
          {
72
            id: 'b0daca4a-ffd8-4865-926b-e24800af2a2d',
73
            values: [0.2331, 1.0125, 0.6131, 0.9421, 0.9661, 0.8121],
74
            metadata: { text: 'She sells seashells by the seashore' },
75
            score: 0.71151,
76
          },
77
          {
78
            id: 'a44706aa-a366-48bc-8cc1-3feffd87d548',
79
            values: [0.2321, 0.8121, 0.6315, 0.6151, 0.4121, 0.1512],
80
            metadata: {
81
              text: 'Peter Piper picked a peck of pickled peppers',
82
            },
83
            score: 0.68913,
84
          },
85
          {
86
            id: '43cfcb31-07e2-411f-8bf9-f82a95ba8b96',
87
            values: [0.0515, 0.7512, 0.8612, 0.2153, 0.15121, 0.6812],
88
            metadata: {
89
              text: 'You know New York, you need New York, you know you need unique New York',
90
            },
91
            score: 0.94812,
92
          },
93
        ],
94
        count: 3,
95
      };
96
      assert.deepStrictEqual(results, expected);
97
    }
98
 
99
    {
100
      // with returnValues = unset (false), returnMetadata = false ("none")
101
      const results = await IDX.query(new Float32Array(5), {
102
        topK: 3,
103
        returnMetadata: false,
104
      });
105
      assert.equal(true, results.count > 0);
106
      /** @type {VectorizeMatches}  */
107
      const expected = {
108
        matches: [
109
          {
110
            id: 'b0daca4a-ffd8-4865-926b-e24800af2a2d',
111
            score: 0.71151,
112
          },
113
          {
114
            id: 'a44706aa-a366-48bc-8cc1-3feffd87d548',
115
            score: 0.68913,
116
          },
117
          {
118
            id: '43cfcb31-07e2-411f-8bf9-f82a95ba8b96',
119
            score: 0.94812,
120
          },
121
        ],
122
        count: 3,
123
      };
124
      assert.deepStrictEqual(results, expected);
125
    }
126
 
127
    {
128
      // with returnValues = unset (false), returnMetadata = true ("all")
129
      const results = await IDX.query(new Float32Array(5), {
130
        topK: 3,
131
        returnMetadata: true,
132
      });
133
      assert.equal(true, results.count > 0);
134
      /** @type {VectorizeMatches}  */
135
      const expected = {
136
        matches: [
137
          {
138
            id: 'b0daca4a-ffd8-4865-926b-e24800af2a2d',
139
            metadata: { text: 'She sells seashells by the seashore' },
140
            score: 0.71151,
141
          },
142
          {
143
            id: 'a44706aa-a366-48bc-8cc1-3feffd87d548',
144
            metadata: {
145
              text: 'Peter Piper picked a peck of pickled peppers',
146
            },
147
            score: 0.68913,
148
          },
149
          {
150
            id: '43cfcb31-07e2-411f-8bf9-f82a95ba8b96',
151
            metadata: {
152
              text: 'You know New York, you need New York, you know you need unique New York',
153
            },
154
            score: 0.94812,
155
          },
156
        ],
157
        count: 3,
158
      };
159
      assert.deepStrictEqual(results, expected);
160
    }
161
 
162
    {
163
      // with returnValues = unset (false), returnMetadata = unset (none)
164
      const results = await IDX.query(new Float32Array(5), {
165
        topK: 3,
166
      });
167
      assert.equal(true, results.count > 0);
168
      /** @type {VectorizeMatches}  */
169
      const expected = {
170
        matches: [
171
          {
172
            id: 'b0daca4a-ffd8-4865-926b-e24800af2a2d',
173
            score: 0.71151,
174
          },
175
          {
176
            id: 'a44706aa-a366-48bc-8cc1-3feffd87d548',
177
            score: 0.68913,
178
          },
179
          {
180
            id: '43cfcb31-07e2-411f-8bf9-f82a95ba8b96',
181
            score: 0.94812,
182
          },
183
        ],
184
        count: 3,
185
      };
186
      assert.deepStrictEqual(results, expected);
187
    }
188
 
189
    {
190
      // with returnValues = unset (false), returnMetadata = unset (none), filter = "Peter Piper picked a peck of pickled peppers"
191
      const results = await IDX.query(new Float32Array(5), {
192
        topK: 1,
193
        filter: {
194
          text: { $eq: 'Peter Piper picked a peck of pickled peppers' },
195
        },
196
      });
197
      assert.equal(true, results.count > 0);
198
      /** @type {VectorizeMatches}  */
199
      const expected = {
200
        matches: [
201
          {
202
            id: 'a44706aa-a366-48bc-8cc1-3feffd87d548',
203
            score: 0.68913,
204
          },
205
        ],
206
        count: 1,
207
      };
208
      assert.deepStrictEqual(results, expected);
209
    }
210
 
211
    {
212
      // with returnValues = unset (false), returnMetadata = unset (none), filter = "Peter Piper picked a peck of pickled peppers"
213
      const results = await IDX.query(new Float32Array(5), {
214
        topK: 1,
215
        filter: {
216
          text: {
217
            $in: [
218
              'Peter Piper picked a peck of pickled peppers',
219
              'She sells seashells by the seashore',
220
            ],
221
          },
222
        },
223
      });
224
      assert.equal(true, results.count > 0);
225
      /** @type {VectorizeMatches}  */
226
      const expected = {
227
        matches: [
228
          {
229
            id: 'b0daca4a-ffd8-4865-926b-e24800af2a2d',
230
            score: 0.71151,
231
          },
232
          {
233
            id: 'a44706aa-a366-48bc-8cc1-3feffd87d548',
234
            score: 0.68913,
235
          },
236
        ],
237
        count: 2,
238
      };
239
      assert.deepStrictEqual(results, expected);
240
    }
241
 
242
    {
243
      const results = await IDX.query(new Float32Array(5), {
244
        topK: 1,
245
        filter: {
246
          text: {
247
            $gt: 'Peter Piper picked a peck of pickled peppers',
248
            $lt: 'You know New York, you need New York, you know you need unique New York',
249
          },
250
        },
251
      });
252
      assert.equal(true, results.count > 0);
253
      /** @type {VectorizeMatches}  */
254
      const expected = {
255
        matches: [
256
          {
257
            id: 'b0daca4a-ffd8-4865-926b-e24800af2a2d',
258
            score: 0.71151,
259
          },
260
        ],
261
        count: 1,
262
      };
263
      assert.deepStrictEqual(results, expected);
264
    }
265
  },
266
};
267
 
268
export const test_vector_search_vector_insert = {
269
  /**
270
   * @param {unknown} _
271
   * @param {Env} env
272
   */
273
  async test(_, env) {
274
    const IDX = env['vector-search'];
275
    {
276
      /** @type {Array<VectorizeVector>}  */
277
      const newVectors = [
278
        {
279
          id: '15cc795d-93d3-416d-9a2a-36fa6fac73da',
280
          values: new Float32Array(),
281
          metadata: { text: 'He threw three free throws' },
282
        },
283
        {
284
          id: '15cc795d-93d3-416d-9a2a-36fa6fac73da',
285
          values: new Float32Array(),
286
          metadata: { text: 'Which witch is which?' },
287
        },
288
      ];
289
      const results = await IDX.insert(newVectors);
290
      assert.equal(results.mutationId, `total vectors: 5`);
291
    }
292
  },
293
};
294
 
295
export const test_vector_search_vector_insert_error = {
296
  /**
297
   * @param {unknown} _
298
   * @param {Env} env
299
   */
300
  async test(_, env) {
301
    const IDX = env['vector-search'];
302
    {
303
      /** @type {Array<VectorizeVector>}  */
304
      const newVectors = [
305
        {
306
          id: 'fail-with-test-error',
307
          values: new Float32Array(),
308
        },
309
      ];
310
 
311
      /** @type {Error | null} */
312
      let error = null;
313
      try {
314
        await IDX.insert(newVectors);
315
      } catch (e) {
316
        error = /** @type {Error} */ (e);
317
      }
318
 
319
      assert.equal(
320
        error && error.message,
321
        'VECTOR_INSERT_ERROR (code = 9999): You asked me for this error'
322
      );
323
    }
324
  },
325
};
326
 
327
export const test_vector_search_vector_upsert = {
328
  /**
329
   * @param {unknown} _
330
   * @param {Env} env
331
   */
332
  async test(_, env) {
333
    const IDX = env['vector-search'];
334
    {
335
      /** @type {Array<VectorizeVector>}  */
336
      const newVectors = [
337
        {
338
          id: '15cc795d-93d3-416d-9a2a-36fa6fac73da',
339
          values: new Float32Array(),
340
          metadata: { text: 'He threw three free throws' },
341
        },
342
        {
343
          id: '15cc795d-93d3-416d-9a2a-36fa6fac73da',
344
          values: [0.3611, 0.9481, 0.8121, 0.7121, 0.8121, 0.0512],
345
          metadata: { text: 'Which witch is which?' },
346
        },
347
      ];
348
      const results = await IDX.upsert(newVectors);
349
      assert.equal(results.mutationId, `total vectors: 4`);
350
    }
351
  },
352
};
353
 
354
export const test_vector_search_vector_delete_ids = {
355
  /**
356
   * @param {unknown} _
357
   * @param {Env} env
358
   */
359
  async test(_, env) {
360
    const IDX = env['vector-search'];
361
    {
362
      const results = await IDX.deleteByIds([
363
        'vector-a',
364
        'vector-b',
365
        'vector-c',
366
      ]);
367
      assert.equal(results.mutationId, `deleted vectors: 3`);
368
    }
369
  },
370
};
371
 
372
export const test_vector_search_vector_get_ids = {
373
  /**
374
   * @param {unknown} _
375
   * @param {Env} env
376
   */
377
  async test(_, env) {
378
    const IDX = env['vector-search'];
379
    {
380
      const results = await IDX.getByIds([
381
        'b0daca4a-ffd8-4865-926b-e24800af2a2d',
382
        '43cfcb31-07e2-411f-8bf9-f82a95ba8b96',
383
      ]);
384
      assert.deepStrictEqual(results, [
385
        {
386
          id: 'b0daca4a-ffd8-4865-926b-e24800af2a2d',
387
          values: [0.2331, 1.0125, 0.6131, 0.9421, 0.9661, 0.8121],
388
          metadata: { text: 'She sells seashells by the seashore' },
389
        },
390
        {
391
          id: '43cfcb31-07e2-411f-8bf9-f82a95ba8b96',
392
          values: [0.0515, 0.7512, 0.8612, 0.2153, 0.15121, 0.6812],
393
          metadata: {
394
            text: 'You know New York, you need New York, you know you need unique New York',
395
          },
396
        },
397
      ]);
398
    }
399
  },
400
};
401
 
402
export const test_vector_search_can_use_enum_exports = {
403
  async test() {
404
    assert.equal(
405
      KnownModel['openai/text-embedding-ada-002'],
406
      'openai/text-embedding-ada-002'
407
    );
408
    assert.equal(DistanceMetric.COSINE, 'cosine');
409
  },
410
};