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
const queryMetadataOptional =
6
  !!Cloudflare.compatibilityFlags['vectorize_query_metadata_optional'];
7
 
8
interface Fetcher {
9
  fetch: typeof fetch;
10
}
11
 
12
const Operation = {
13
  INDEX_GET: 'INDEX_GET',
14
  VECTOR_QUERY: 'VECTOR_QUERY',
15
  VECTOR_INSERT: 'VECTOR_INSERT',
16
  VECTOR_UPSERT: 'VECTOR_UPSERT',
17
  VECTOR_GET: 'VECTOR_GET',
18
  VECTOR_DELETE: 'VECTOR_DELETE',
19
} as const;
20
type OperationKey = keyof typeof Operation;
21
 
22
type VectorizeVersion = 'v1' | 'v2';
23
 
24
type QueryImplV2Params =
25
  | { vector: VectorFloatArray | number[]; vectorId?: undefined }
26
  | { vector?: undefined; vectorId: string };
27
 
28
function toNdJson(arr: object[]): string {
29
  return arr.reduce((acc, o) => acc + JSON.stringify(o) + '\n', '').trim();
30
}
31
 
32
/*
33
 * The Vectorize beta VectorizeIndex shares the same methods, so to keep things simple, they share one implementation.
34
 * The types here are specific to Vectorize GA, but the types here don't actually matter as they are stripped away
35
 * and not visible to end users.
36
 */
37
class VectorizeIndexImpl implements Vectorize {
38
  // eslint-disable-next-line no-restricted-syntax
39
  private readonly fetcher: Fetcher;
40
  // eslint-disable-next-line no-restricted-syntax
41
  private readonly indexId: string;
42
  // eslint-disable-next-line no-restricted-syntax
43
  private readonly indexVersion: VectorizeVersion;
44
  // eslint-disable-next-line no-restricted-syntax
45
  private readonly useNdJson: boolean;
46
 
47
  constructor(
48
    fetcher: Fetcher,
49
    indexId: string,
50
    indexVersion: VectorizeVersion,
51
    useNdJson: boolean
52
  ) {
53
    this.fetcher = fetcher;
54
    this.indexId = indexId;
55
    this.indexVersion = indexVersion;
56
    this.useNdJson = useNdJson;
57
  }
58
 
59
  async describe(): Promise<VectorizeIndexInfo> {
60
    const endpoint =
61
      this.indexVersion === 'v2' ? `info` : `binding/indexes/${this.indexId}`;
62
    const res = await this._send(Operation.INDEX_GET, endpoint, {
63
      method: 'GET',
64
    });
65
 
66
    return await toJson<VectorizeIndexInfo>(res);
67
  }
68
 
69
  async query(
70
    vector: VectorFloatArray | number[],
71
    options?: VectorizeQueryOptions
72
  ): Promise<VectorizeMatches> {
73
    if (this.indexVersion === 'v2') {
74
      return await this.queryImplV2(
75
        { vector: Array.isArray(vector) ? vector : Array.from(vector) },
76
        options
77
      );
78
    } else {
79
      if (
80
        options &&
81
        options.returnMetadata &&
82
        typeof options.returnMetadata !== 'boolean'
83
      ) {
84
        throw new Error(
85
          `Invalid returnMetadata option. Expected boolean; got: ${options.returnMetadata}`
86
        );
87
      }
88
      const compat = {
89
        queryMetadataOptional,
90
      };
91
      const res = await this._send(
92
        Operation.VECTOR_QUERY,
93
        `binding/indexes/${this.indexId}/query`,
94
        {
95
          method: 'POST',
96
          body: JSON.stringify({
97
            ...options,
98
            vector: Array.isArray(vector) ? vector : Array.from(vector),
99
            compat,
100
          }),
101
          headers: {
102
            'content-type': 'application/json',
103
            accept: 'application/json',
104
            'cf-vector-search-query-compat': JSON.stringify(compat),
105
          },
106
        }
107
      );
108
 
109
      return await toJson<VectorizeMatches>(res);
110
    }
111
  }
112
 
113
  async queryById(
114
    vectorId: string,
115
    options?: VectorizeQueryOptions
116
  ): Promise<VectorizeMatches> {
117
    if (this.indexVersion === 'v1') {
118
      throw new Error(`QueryById operation is not supported for v1 indexes.`);
119
    } else {
120
      return await this.queryImplV2({ vectorId }, options);
121
    }
122
  }
123
 
124
  async insert(vectors: VectorizeVector[]): Promise<VectorizeAsyncMutation> {
125
    const endpoint =
126
      this.indexVersion === 'v2'
127
        ? `insert`
128
        : `binding/indexes/${this.indexId}/insert`;
129
    const bodyVecArr = vectors.map((vec) => ({
130
      ...vec,
131
      values: Array.isArray(vec.values) ? vec.values : Array.from(vec.values),
132
    }));
133
 
134
    const body = this.useNdJson
135
      ? toNdJson(bodyVecArr)
136
      : JSON.stringify({ vectors: bodyVecArr });
137
 
138
    const contentType = this.useNdJson
139
      ? 'application/x-ndjson'
140
      : 'application/json';
141
 
142
    const res = await this._send(Operation.VECTOR_INSERT, endpoint, {
143
      method: 'POST',
144
      body,
145
      headers: {
146
        'content-type': contentType,
147
        'cf-vector-search-dim-width': String(
148
          // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
149
          vectors.length ? vectors.at(0)?.values?.length : 0
150
        ),
151
        'cf-vector-search-dim-height': String(vectors.length),
152
        accept: 'application/json',
153
      },
154
    });
155
 
156
    return await toJson<VectorizeAsyncMutation>(res);
157
  }
158
 
159
  async upsert(vectors: VectorizeVector[]): Promise<VectorizeAsyncMutation> {
160
    const endpoint =
161
      this.indexVersion === 'v2'
162
        ? `upsert`
163
        : `binding/indexes/${this.indexId}/upsert`;
164
    const bodyVecArr = vectors.map((vec) => ({
165
      ...vec,
166
      values: Array.isArray(vec.values) ? vec.values : Array.from(vec.values),
167
    }));
168
 
169
    const body = this.useNdJson
170
      ? toNdJson(bodyVecArr)
171
      : JSON.stringify({ vectors: bodyVecArr });
172
 
173
    const contentType = this.useNdJson
174
      ? 'application/x-ndjson'
175
      : 'application/json';
176
 
177
    const res = await this._send(Operation.VECTOR_UPSERT, endpoint, {
178
      method: 'POST',
179
      body,
180
      headers: {
181
        'content-type': contentType,
182
        'cf-vector-search-dim-width': String(
183
          // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
184
          vectors.length ? vectors.at(0)?.values?.length : 0
185
        ),
186
        'cf-vector-search-dim-height': String(vectors.length),
187
        accept: 'application/json',
188
      },
189
    });
190
 
191
    return await toJson<VectorizeAsyncMutation>(res);
192
  }
193
 
194
  async getByIds(ids: string[]): Promise<VectorizeVector[]> {
195
    const endpoint =
196
      this.indexVersion === 'v2'
197
        ? `getByIds`
198
        : `binding/indexes/${this.indexId}/getByIds`;
199
    const res = await this._send(Operation.VECTOR_GET, endpoint, {
200
      method: 'POST',
201
      body: JSON.stringify({ ids }),
202
      headers: {
203
        'content-type': 'application/json',
204
        accept: 'application/json',
205
      },
206
    });
207
 
208
    return await toJson<VectorizeVector[]>(res);
209
  }
210
 
211
  async deleteByIds(ids: string[]): Promise<VectorizeAsyncMutation> {
212
    const endpoint =
213
      this.indexVersion === 'v2'
214
        ? `deleteByIds`
215
        : `binding/indexes/${this.indexId}/deleteByIds`;
216
    const res = await this._send(Operation.VECTOR_DELETE, endpoint, {
217
      method: 'POST',
218
      body: JSON.stringify({ ids }),
219
      headers: {
220
        'content-type': 'application/json',
221
        accept: 'application/json',
222
      },
223
    });
224
 
225
    return await toJson<VectorizeAsyncMutation>(res);
226
  }
227
 
228
  // eslint-disable-next-line no-restricted-syntax
229
  private async _send(
230
    operation: OperationKey,
231
    endpoint: string,
232
    init: RequestInit
233
  ): Promise<Response> {
234
    const res = await this.fetcher.fetch(
235
      `http://vector-search/${endpoint}`, // `http://vector-search` is just a dummy host, the attached fetcher will receive the request
236
      init
237
    );
238
    if (res.status !== 200) {
239
      let err: Error | null = null;
240
 
241
      try {
242
        const errResponse = (await res.json()) as VectorizeError;
243
        err = new Error(
244
          `${Operation[operation]}_ERROR${
245
            typeof errResponse.code === 'number'
246
              ? ` (code = ${errResponse.code})`
247
              : ''
248
          }: ${errResponse.error}`,
249
          {
250
            cause: new Error(errResponse.error),
251
          }
252
        );
253
      } catch {
254
        // do nothing
255
      }
256
 
257
      if (err) {
258
        throw err;
259
      } else {
260
        throw new Error(
261
          `${Operation[operation]}_ERROR: Status + ${res.status}`,
262
          {
263
            cause: new Error(`Status ${res.status}`),
264
          }
265
        );
266
      }
267
    }
268
 
269
    return res;
270
  }
271
 
272
  // eslint-disable-next-line no-restricted-syntax
273
  private async queryImplV2(
274
    vectorParams: QueryImplV2Params,
275
    options?: VectorizeQueryOptions
276
  ): Promise<VectorizeMatches> {
277
    if (options?.returnMetadata) {
278
      if (
279
        typeof options.returnMetadata !== 'boolean' &&
280
        !isVectorizeMetadataRetrievalLevel(options.returnMetadata)
281
      ) {
282
        throw new Error(
283
          `Invalid returnMetadata option. Expected: true, false, "none", "indexed" or "all"; got: ${options.returnMetadata}`
284
        );
285
      }
286
 
287
      if (typeof options.returnMetadata === 'boolean') {
288
        // Allow boolean returnMetadata for backward compatibility. true converts to 'all' and false converts to 'none'
289
        // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
290
        options.returnMetadata = options.returnMetadata ? 'all' : 'none';
291
      }
292
    }
293
    const res = await this._send(Operation.VECTOR_QUERY, `query`, {
294
      method: 'POST',
295
      body: JSON.stringify({
296
        ...options,
297
        ...(vectorParams.vector
298
          ? { vector: vectorParams.vector }
299
          : { vectorId: vectorParams.vectorId }),
300
      }),
301
      headers: {
302
        'content-type': 'application/json',
303
        accept: 'application/json',
304
      },
305
    });
306
 
307
    return await toJson<VectorizeMatches>(res);
308
  }
309
}
310
 
311
function isVectorizeMetadataRetrievalLevel(value: unknown): boolean {
312
  return (
313
    typeof value === 'string' &&
314
    (value === 'all' || value === 'indexed' || value === 'none')
315
  );
316
}
317
 
318
const maxBodyLogChars = 1_000;
319
async function toJson<T = unknown>(response: Response): Promise<T> {
320
  const body = await response.text();
321
  try {
322
    return JSON.parse(body) as T;
323
  } catch {
324
    throw new Error(
325
      `Failed to parse body as JSON, got: ${
326
        body.length > maxBodyLogChars
327
          ? `${body.slice(0, maxBodyLogChars)}…`
328
          : body
329
      }`
330
    );
331
  }
332
}
333
 
334
export function makeBinding(env: {
335
  fetcher: Fetcher;
336
  indexId: string;
337
  indexVersion?: VectorizeVersion;
338
  useNdJson?: boolean;
339
}): Vectorize {
340
  return new VectorizeIndexImpl(
341
    env.fetcher,
342
    env.indexId,
343
    env.indexVersion ?? 'v1',
344
    env.useNdJson ?? false
345
  );
346
}
347
 
348
export default makeBinding;