triangle-utils 1.4.88 → 1.4.89

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 (6) hide show
  1. package/dist/src/UtilsS3.js +31 -31
  2. package/dist/src/UtilsS3Vectors.d.ts +17 -12
  3. package/dist/src/UtilsS3Vectors.js +150 -67
  4. package/package.json +1 -1
  5. package/src/UtilsS3.ts +32 -32
  6. package/src/UtilsS3Vectors.ts +160 -78
package/dist/src/UtilsS3.js CHANGED
@@ -1,7 +1,7 @@
1
1
  import { S3, NoSuchKey, GetObjectCommand, PutObjectCommand } from "@aws-sdk/client-s3";
2
2
  import { getSignedUrl } from "@aws-sdk/s3-request-presigner";
3
3
  import { Readable } from "stream";
4
- function get_s3_key(s3_id) {
4
+ function get_s3_input(s3_id) {
5
5
  const [bucket, path] = (s3_id.match(/^s3\:\/\/([^\/]+)\/([^$]+)$/) || []).slice(1, 3);
6
6
  if (bucket === undefined || path === undefined) {
7
7
  return undefined;
@@ -14,12 +14,12 @@ export class UtilsS3 {
14
14
  this.s3 = new S3({ region: region });
15
15
  }
16
16
  async get_headers(s3_id) {
17
- const s3_key = get_s3_key(s3_id);
18
- if (s3_key === undefined) {
17
+ const s3_input = get_s3_input(s3_id);
18
+ if (s3_input === undefined) {
19
19
  return undefined;
20
20
  }
21
21
  try {
22
- const headers = await this.s3.headObject(s3_key);
22
+ const headers = await this.s3.headObject(s3_input);
23
23
  return {
24
24
  size: headers.ContentLength,
25
25
  ...headers
@@ -30,12 +30,12 @@ export class UtilsS3 {
30
30
  }
31
31
  }
32
32
  async get(s3_id, encoding = "utf-8") {
33
- const s3_key = get_s3_key(s3_id);
34
- if (s3_key === undefined) {
33
+ const s3_input = get_s3_input(s3_id);
34
+ if (s3_input === undefined) {
35
35
  return undefined;
36
36
  }
37
37
  try {
38
- const response = await this.s3.getObject(s3_key);
38
+ const response = await this.s3.getObject(s3_input);
39
39
  const body = response.Body;
40
40
  if (body === undefined) {
41
41
  return undefined;
@@ -51,12 +51,12 @@ export class UtilsS3 {
51
51
  }
52
52
  }
53
53
  async get_buffer(s3_id) {
54
- const s3_key = get_s3_key(s3_id);
55
- if (s3_key === undefined) {
54
+ const s3_input = get_s3_input(s3_id);
55
+ if (s3_input === undefined) {
56
56
  return undefined;
57
57
  }
58
58
  try {
59
- const response = await this.s3.getObject(s3_key);
59
+ const response = await this.s3.getObject(s3_input);
60
60
  const body = response.Body;
61
61
  if (body === undefined) {
62
62
  return undefined;
@@ -72,12 +72,12 @@ export class UtilsS3 {
72
72
  }
73
73
  }
74
74
  async get_stream(s3_id) {
75
- const s3_key = get_s3_key(s3_id);
76
- if (s3_key === undefined) {
75
+ const s3_input = get_s3_input(s3_id);
76
+ if (s3_input === undefined) {
77
77
  return undefined;
78
78
  }
79
79
  try {
80
- const response = await this.s3.getObject(s3_key);
80
+ const response = await this.s3.getObject(s3_input);
81
81
  const body = response.Body;
82
82
  if (body === undefined) {
83
83
  return undefined;
@@ -96,17 +96,17 @@ export class UtilsS3 {
96
96
  }
97
97
  async query_prefix(s3_id_prefix, options = {}) {
98
98
  const compile = options.compile !== undefined ? options.compile : false;
99
- const s3_key = get_s3_key(s3_id_prefix);
100
- if (s3_key === undefined) {
99
+ const s3_input = get_s3_input(s3_id_prefix);
100
+ if (s3_input === undefined) {
101
101
  return [];
102
102
  }
103
103
  const s3_ids = [];
104
104
  let last_token = undefined;
105
105
  while (true) {
106
106
  const request = {
107
- Bucket: s3_key.Bucket,
107
+ Bucket: s3_input.Bucket,
108
108
  Delimiter: "/",
109
- Prefix: s3_key.Key,
109
+ Prefix: s3_input.Key,
110
110
  ContinuationToken: last_token
111
111
  };
112
112
  const response = await this.s3.listObjectsV2(request);
@@ -114,12 +114,12 @@ export class UtilsS3 {
114
114
  .map(content => content.Key)
115
115
  .filter(key => key !== undefined)
116
116
  .filter(key => key.substring(key.length - 1) !== "/")
117
- .map(key => "s3://" + s3_key.Bucket + "/" + key);
117
+ .map(key => "s3://" + s3_input.Bucket + "/" + key);
118
118
  s3_ids.push(...new_s3_ids);
119
119
  const sub_s3_id_prefixes = (response.CommonPrefixes || [])
120
120
  .map(prefix => prefix.Prefix)
121
121
  .filter(prefix => prefix !== undefined)
122
- .map(prefix => "s3://" + s3_key.Bucket + "/" + prefix);
122
+ .map(prefix => "s3://" + s3_input.Bucket + "/" + prefix);
123
123
  for (const sub_s3_id_prefix of sub_s3_id_prefixes) {
124
124
  const new_sub_s3_ids = await this.query_prefix(sub_s3_id_prefix);
125
125
  s3_ids.push(...new_sub_s3_ids);
@@ -136,42 +136,42 @@ export class UtilsS3 {
136
136
  if (headers !== undefined) {
137
137
  return undefined;
138
138
  }
139
- const s3_key = get_s3_key(s3_id);
140
- if (s3_key === undefined) {
139
+ const s3_input = get_s3_input(s3_id);
140
+ if (s3_input === undefined) {
141
141
  return undefined;
142
142
  }
143
143
  return await this.s3.putObject({
144
- ...s3_key,
144
+ ...s3_input,
145
145
  Body: content,
146
146
  ContentType: content_type
147
147
  });
148
148
  }
149
149
  async delete(s3_id) {
150
- const s3_key = get_s3_key(s3_id);
151
- if (s3_key === undefined) {
150
+ const s3_input = get_s3_input(s3_id);
151
+ if (s3_input === undefined) {
152
152
  return undefined;
153
153
  }
154
- await this.s3.deleteObject(s3_key);
154
+ await this.s3.deleteObject(s3_input);
155
155
  }
156
156
  async get_download_url(s3_id) {
157
157
  const headers = await this.get_headers(s3_id);
158
158
  if (headers === undefined) {
159
159
  return undefined;
160
160
  }
161
- const s3_key = get_s3_key(s3_id);
162
- if (s3_key === undefined) {
161
+ const s3_input = get_s3_input(s3_id);
162
+ if (s3_input === undefined) {
163
163
  return undefined;
164
164
  }
165
- const command = new GetObjectCommand(s3_key);
165
+ const command = new GetObjectCommand(s3_input);
166
166
  const url = await getSignedUrl(this.s3, command);
167
167
  return url;
168
168
  }
169
169
  async get_upload_url(s3_id) {
170
- const s3_key = get_s3_key(s3_id);
171
- if (s3_key === undefined) {
170
+ const s3_input = get_s3_input(s3_id);
171
+ if (s3_input === undefined) {
172
172
  return undefined;
173
173
  }
174
- const command = new PutObjectCommand(s3_key);
174
+ const command = new PutObjectCommand(s3_input);
175
175
  const url = await getSignedUrl(this.s3, command);
176
176
  return url;
177
177
  }
package/dist/src/UtilsS3Vectors.d.ts CHANGED
@@ -1,17 +1,22 @@
1
- export interface Vector {
2
- key: string;
3
- data: number[];
4
- metadata: any;
1
+ export declare class S3Vector {
2
+ readonly s3_vector_id: string;
3
+ readonly coordinates: number[];
4
+ readonly headers: Record<string, any>;
5
+ constructor(s3_vector: S3Vector);
6
+ static is(s3_vector: any): s3_vector is S3Vector;
5
7
  }
6
8
  export declare class UtilsS3Vectors {
7
9
  private readonly s3vectors;
8
10
  constructor(region: string);
9
- scan(vector_bucket: string, vector_index: string): Promise<string[]>;
10
- get(vector_bucket: string, vector_index: string, vector_key: string): Promise<Vector | undefined>;
11
- batch_get(vector_bucket: string, vector_index: string, vector_keys: string[]): Promise<(Vector | undefined)[]>;
12
- create(vector_bucket: string, vector_index: string, vector: Vector): Promise<void>;
13
- batch_create(vector_bucket: string, vector_index: string, vectors: Vector[]): Promise<void>;
14
- delete(vector_bucket: string, vector_index: string, vector_key: string): Promise<void>;
15
- batch_delete(vector_bucket: string, vector_index: string, vector_keys: string[]): Promise<void>;
16
- query(vector_bucket: string, vector_index: string, data: number[], num_vectors: number, filters?: Record<string, any> | undefined): Promise<import("@aws-sdk/client-s3vectors").QueryOutputVector[] | undefined>;
11
+ scan(s3_vector_id_prefix: string): Promise<string[]>;
12
+ get(s3_vector_id: string): Promise<S3Vector | undefined>;
13
+ batch_get(s3_vector_ids: string[]): Promise<(S3Vector | undefined)[]>;
14
+ put(s3_vector: S3Vector): Promise<void>;
15
+ batch_put(s3_vectors: S3Vector[]): Promise<void>;
16
+ delete(s3_vector_id: string): Promise<void>;
17
+ batch_delete(s3_vector_ids: string[]): Promise<void>;
18
+ query(s3_vector_id_prefix: string, coordinates: number[], num_vectors: number, filters?: Record<string, any> | undefined): Promise<{
19
+ s3_vector_id: string;
20
+ distance: number;
21
+ }[]>;
17
22
  }
package/dist/src/UtilsS3Vectors.js CHANGED
@@ -1,23 +1,47 @@
1
1
  import { S3Vectors } from "@aws-sdk/client-s3vectors";
2
- function get_vector_key(vector_id) {
3
- const [vector_bucket_name, vector_index_name, vector_key] = (vector_id.match(/^s3vectors\:\/\/([^\/]+)\/([^\/]+)\/([^$]+)$/) || []).slice(1, 4);
4
- if (vector_bucket_name === undefined || vector_index_name === undefined || vector_key === undefined) {
5
- return undefined;
2
+ export class S3Vector {
3
+ s3_vector_id;
4
+ coordinates;
5
+ headers;
6
+ constructor(s3_vector) {
7
+ if (!S3Vector.is(s3_vector)) {
8
+ throw Error("Invalid input.");
9
+ }
10
+ this.s3_vector_id = s3_vector.s3_vector_id;
11
+ this.coordinates = s3_vector.coordinates;
12
+ this.headers = s3_vector.headers;
13
+ }
14
+ static is(s3_vector) {
15
+ return (s3_vector !== undefined &&
16
+ typeof s3_vector.s3_vector_id === "string" &&
17
+ Array.isArray(s3_vector.coordinates) && s3_vector.coordinates.every((coordinate) => typeof coordinate === "number") &&
18
+ typeof s3_vector.headers === "object" && Object.keys(s3_vector.headers).every((header_id) => typeof header_id === "string"));
6
19
  }
7
- return { vector_bucket_name: vector_bucket_name, vector_index_name: vector_index_name, vector_key: vector_key };
20
+ }
21
+ function get_s3vectors_input(s3_vector_id) {
22
+ const [s3_vector_bucket_name, s3_vector_index_name, s3_vector_key] = (s3_vector_id.match(/^s3vectors:\/\/([^\/]+)\/([^\/]+)\/([^$]*)$/) || []).slice(1, 4);
23
+ if (s3_vector_bucket_name === undefined || s3_vector_index_name === undefined || s3_vector_key === undefined) {
24
+ throw Error("Bad s3_vector_id: " + s3_vector_id);
25
+ }
26
+ return {
27
+ s3_vector_bucket_name: s3_vector_bucket_name,
28
+ s3_vector_index_name: s3_vector_index_name,
29
+ s3_vector_key: s3_vector_key
30
+ };
8
31
  }
9
32
  export class UtilsS3Vectors {
10
33
  s3vectors;
11
34
  constructor(region) {
12
35
  this.s3vectors = new S3Vectors({ region: region });
13
36
  }
14
- async scan(vector_bucket, vector_index) {
15
- const vector_keys = [];
37
+ async scan(s3_vector_id_prefix) {
38
+ const { s3_vector_bucket_name, s3_vector_index_name } = get_s3vectors_input(s3_vector_id_prefix);
39
+ const s3_vector_ids = [];
16
40
  let next_token = undefined;
17
41
  while (true) {
18
42
  const request = {
19
- vectorBucketName: vector_bucket,
20
- indexName: vector_index,
43
+ vectorBucketName: s3_vector_bucket_name,
44
+ indexName: s3_vector_index_name,
21
45
  maxResults: 1000,
22
46
  nextToken: next_token
23
47
  };
@@ -25,38 +49,54 @@ export class UtilsS3Vectors {
25
49
  if (response === undefined || response.vectors === undefined) {
26
50
  return [];
27
51
  }
28
- vector_keys.push(...response.vectors.map(vector => vector.key));
52
+ s3_vector_ids.push(...response.vectors
53
+ .map(vector => vector.key)
54
+ .filter(s3_vector_key => s3_vector_key !== undefined)
55
+ .map(s3_vector_key => "s3vectors://" + s3_vector_bucket_name + "/" + s3_vector_index_name + "/" + s3_vector_key));
29
56
  next_token = response.nextToken;
30
57
  if (next_token === undefined) {
31
- return vector_keys.filter(vector_key => vector_key !== undefined);
58
+ return s3_vector_ids;
32
59
  }
33
60
  }
34
61
  }
35
- async get(vector_bucket, vector_index, vector_key) {
62
+ async get(s3_vector_id) {
63
+ const { s3_vector_bucket_name, s3_vector_index_name, s3_vector_key } = get_s3vectors_input(s3_vector_id);
36
64
  const response = await this.s3vectors.getVectors({
37
- vectorBucketName: vector_bucket,
38
- indexName: vector_index,
39
- keys: [vector_key],
65
+ vectorBucketName: s3_vector_bucket_name,
66
+ indexName: s3_vector_index_name,
67
+ keys: [s3_vector_key],
40
68
  returnData: true,
41
69
  returnMetadata: true
42
70
  });
43
71
  const vector = response.vectors?.[0];
44
- if (vector === undefined || vector.key === undefined || vector.data === undefined || vector.data.float32 === undefined || vector.metadata === undefined) {
72
+ if (vector === undefined) {
45
73
  return undefined;
46
74
  }
47
- return {
48
- key: vector.key,
49
- data: vector.data.float32,
50
- metadata: vector.metadata
75
+ const coordinates = vector.data?.float32;
76
+ const headers = vector.metadata;
77
+ const s3_vector = {
78
+ s3_vector_id: s3_vector_id,
79
+ coordinates: coordinates,
80
+ headers: headers
51
81
  };
82
+ if (!S3Vector.is(s3_vector)) {
83
+ console.log("Bad S3Vector:", s3_vector);
84
+ return undefined;
85
+ }
86
+ return s3_vector;
52
87
  }
53
- async batch_get(vector_bucket, vector_index, vector_keys) {
54
- const vectors = [];
55
- for (let i = 0; i < vector_keys.length; i += 100) {
88
+ async batch_get(s3_vector_ids) {
89
+ const s3vectors_inputs = s3_vector_ids.map(get_s3vectors_input);
90
+ const { s3_vector_bucket_name, s3_vector_index_name } = s3vectors_inputs[0];
91
+ if (!s3vectors_inputs.every(s3vectors_input => s3vectors_input.s3_vector_bucket_name === s3_vector_bucket_name && s3vectors_input.s3_vector_index_name === s3_vector_index_name)) {
92
+ return [];
93
+ }
94
+ const s3_vectors = [];
95
+ for (let i = 0; i < s3vectors_inputs.length; i += 100) {
56
96
  const new_vectors = await this.s3vectors.getVectors({
57
- vectorBucketName: vector_bucket,
58
- indexName: vector_index,
59
- keys: vector_keys.slice(i, i + 100),
97
+ vectorBucketName: s3_vector_bucket_name,
98
+ indexName: s3_vector_index_name,
99
+ keys: s3vectors_inputs.slice(i, i + 100).map(s3vectors_input => s3vectors_input.s3_vector_key),
60
100
  returnData: true,
61
101
  returnMetadata: true
62
102
  })
@@ -64,75 +104,118 @@ export class UtilsS3Vectors {
64
104
  if (new_vectors === undefined) {
65
105
  return [];
66
106
  }
67
- const vectors_by_key = Object.fromEntries(new_vectors
68
- .map(vector => [
69
- vector.key,
70
- vector.key !== undefined && vector.data !== undefined && vector.data.float32 !== undefined && vector.metadata !== undefined ? {
71
- key: vector.key,
72
- data: vector.data.float32,
73
- metadata: vector.metadata
74
- } : undefined
75
- ]));
76
- vectors.push(...vector_keys.map(vector_key => vectors_by_key[vector_key]));
107
+ const new_s3_vectors = [];
108
+ for (const new_vector of new_vectors) {
109
+ const s3_vector_key = new_vector.key;
110
+ const coordinates = new_vector.data?.float32;
111
+ const headers = new_vector.metadata;
112
+ const new_s3_vector = {
113
+ s3_vector_id: "s3vectors://" + s3_vector_bucket_name + "/" + s3_vector_index_name + "/" + s3_vector_key,
114
+ coordinates: coordinates,
115
+ headers: headers
116
+ };
117
+ if (!S3Vector.is(new_s3_vector)) {
118
+ console.log("Bad S3Vector:", new_s3_vector);
119
+ continue;
120
+ }
121
+ new_s3_vectors.push(new_s3_vector);
122
+ }
123
+ const new_s3_vectors_by_s3_vector_id = Object.fromEntries(new_s3_vectors.map(new_s3_vector => [new_s3_vector.s3_vector_id, new_s3_vector]));
124
+ s3_vectors.push(...s3_vector_ids.slice(i, i + 100).map(s3_vector_id => new_s3_vectors_by_s3_vector_id[s3_vector_id]));
77
125
  }
78
- return vectors;
126
+ return s3_vectors;
79
127
  }
80
- async create(vector_bucket, vector_index, vector) {
128
+ async put(s3_vector) {
129
+ const { s3_vector_bucket_name, s3_vector_index_name, s3_vector_key } = get_s3vectors_input(s3_vector.s3_vector_id);
81
130
  await this.s3vectors.putVectors({
82
- vectorBucketName: vector_bucket,
83
- indexName: vector_index,
131
+ vectorBucketName: s3_vector_bucket_name,
132
+ indexName: s3_vector_index_name,
84
133
  vectors: [{
85
- key: vector.key,
134
+ key: s3_vector_key,
86
135
  data: {
87
- float32: vector.data
136
+ float32: s3_vector.coordinates
88
137
  },
89
- metadata: vector.metadata
138
+ metadata: s3_vector.headers
90
139
  }]
91
140
  });
92
141
  }
93
- async batch_create(vector_bucket, vector_index, vectors) {
94
- for (let i = 0; i < vectors.length; i += 100) {
142
+ async batch_put(s3_vectors) {
143
+ const s3vectors_inputs = s3_vectors.map(s3_vector => get_s3vectors_input(s3_vector.s3_vector_id));
144
+ const { s3_vector_bucket_name, s3_vector_index_name } = s3vectors_inputs[0];
145
+ if (!s3vectors_inputs.every(s3vectors_input => s3vectors_input.s3_vector_bucket_name === s3_vector_bucket_name && s3vectors_input.s3_vector_index_name === s3_vector_index_name)) {
146
+ return undefined;
147
+ }
148
+ for (let i = 0; i < s3_vectors.length; i += 100) {
95
149
  await this.s3vectors.putVectors({
96
- vectorBucketName: vector_bucket,
97
- indexName: vector_index,
98
- vectors: vectors.slice(i, i + 100).map(vector => ({
99
- key: vector.key,
150
+ vectorBucketName: s3_vector_bucket_name,
151
+ indexName: s3_vector_index_name,
152
+ vectors: s3_vectors.slice(i, i + 100).map(s3_vector => ({
153
+ key: get_s3vectors_input(s3_vector.s3_vector_id).s3_vector_key,
100
154
  data: {
101
- float32: vector.data
155
+ float32: s3_vector.coordinates
102
156
  },
103
- metadata: vector.metadata
157
+ metadata: s3_vector.headers
104
158
  }))
105
159
  });
106
160
  }
107
161
  }
108
- async delete(vector_bucket, vector_index, vector_key) {
162
+ async delete(s3_vector_id) {
163
+ const { s3_vector_bucket_name, s3_vector_index_name, s3_vector_key } = get_s3vectors_input(s3_vector_id);
109
164
  await this.s3vectors.deleteVectors({
110
- vectorBucketName: vector_bucket,
111
- indexName: vector_index,
112
- keys: [vector_key]
165
+ vectorBucketName: s3_vector_bucket_name,
166
+ indexName: s3_vector_index_name,
167
+ keys: [s3_vector_key]
113
168
  });
114
169
  }
115
- async batch_delete(vector_bucket, vector_index, vector_keys) {
116
- for (let i = 0; i < vector_keys.length; i += 100) {
170
+ async batch_delete(s3_vector_ids) {
171
+ const s3vectors_inputs = s3_vector_ids.map(get_s3vectors_input);
172
+ const { s3_vector_bucket_name, s3_vector_index_name } = s3vectors_inputs[0];
173
+ if (!s3vectors_inputs.every(s3vectors_input => s3vectors_input.s3_vector_bucket_name === s3_vector_bucket_name && s3vectors_input.s3_vector_index_name === s3_vector_index_name)) {
174
+ return undefined;
175
+ }
176
+ for (let i = 0; i < s3vectors_inputs.length; i += 100) {
117
177
  await this.s3vectors.deleteVectors({
118
- vectorBucketName: vector_bucket,
119
- indexName: vector_index,
120
- keys: vector_keys.slice(i, i + 100)
178
+ vectorBucketName: s3_vector_bucket_name,
179
+ indexName: s3_vector_index_name,
180
+ keys: s3vectors_inputs.slice(i, i + 100).map(s3vectors_input => s3vectors_input.s3_vector_key)
121
181
  });
122
182
  }
123
183
  }
124
- async query(vector_bucket, vector_index, data, num_vectors, filters = undefined) {
125
- const response = await this.s3vectors.queryVectors({
126
- vectorBucketName: vector_bucket,
127
- indexName: vector_index,
184
+ async query(s3_vector_id_prefix, coordinates, num_vectors, filters = undefined) {
185
+ const { s3_vector_bucket_name, s3_vector_index_name } = get_s3vectors_input(s3_vector_id_prefix);
186
+ const new_vectors = await this.s3vectors.queryVectors({
187
+ vectorBucketName: s3_vector_bucket_name,
188
+ indexName: s3_vector_index_name,
128
189
  queryVector: {
129
- float32: data
190
+ float32: coordinates
130
191
  },
131
192
  topK: num_vectors,
132
193
  returnDistance: true,
133
194
  returnMetadata: true,
134
195
  filter: filters
135
- });
136
- return response.vectors;
196
+ })
197
+ .then(response => response.vectors);
198
+ if (new_vectors === undefined) {
199
+ return [];
200
+ }
201
+ const results = [];
202
+ for (const new_vector of new_vectors) {
203
+ const s3_vector_key = new_vector.key;
204
+ if (s3_vector_key === undefined) {
205
+ console.log("Bad s3_vector_key:", s3_vector_key);
206
+ continue;
207
+ }
208
+ const s3_vector_id = s3_vector_id_prefix + "/" + s3_vector_key;
209
+ const distance = new_vector.distance;
210
+ if (distance === undefined) {
211
+ console.log("Bad distance:", distance);
212
+ continue;
213
+ }
214
+ results.push({
215
+ s3_vector_id: s3_vector_id,
216
+ distance: distance
217
+ });
218
+ }
219
+ return results;
137
220
  }
138
221
  }
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "triangle-utils",
3
- "version": "1.4.88",
3
+ "version": "1.4.89",
4
4
  "main": "dist/src/index.js",
5
5
  "types": "dist/src/index.d.ts",
6
6
  "directories": {
package/src/UtilsS3.ts CHANGED
@@ -2,12 +2,12 @@ import { S3, NoSuchKey, GetObjectCommand, PutObjectCommand, ListObjectsV2Command
2
2
  import { getSignedUrl } from "@aws-sdk/s3-request-presigner"
3
3
  import { Readable } from "stream"
4
4
 
5
- interface S3Key {
5
+ interface S3Input {
6
6
  Bucket : string,
7
7
  Key : string
8
8
  }
9
9
 
10
- function get_s3_key(s3_id : string) : S3Key | undefined {
10
+ function get_s3_input(s3_id : string) : S3Input | undefined {
11
11
  const [bucket, path] = (s3_id.match(/^s3\:\/\/([^\/]+)\/([^$]+)$/) || []).slice(1, 3)
12
12
  if (bucket === undefined || path === undefined) {
13
13
  return undefined
@@ -24,12 +24,12 @@ export class UtilsS3 {
24
24
  }
25
25
 
26
26
  async get_headers(s3_id : string) {
27
- const s3_key = get_s3_key(s3_id)
28
- if (s3_key === undefined) {
27
+ const s3_input = get_s3_input(s3_id)
28
+ if (s3_input === undefined) {
29
29
  return undefined
30
30
  }
31
31
  try {
32
- const headers = await this.s3.headObject(s3_key)
32
+ const headers = await this.s3.headObject(s3_input)
33
33
  return {
34
34
  size : headers.ContentLength,
35
35
  ...headers
@@ -40,12 +40,12 @@ export class UtilsS3 {
40
40
  }
41
41
 
42
42
  async get(s3_id : string, encoding : string = "utf-8") : Promise<string | undefined> {
43
- const s3_key = get_s3_key(s3_id)
44
- if (s3_key === undefined) {
43
+ const s3_input = get_s3_input(s3_id)
44
+ if (s3_input === undefined) {
45
45
  return undefined
46
46
  }
47
47
  try {
48
- const response = await this.s3.getObject(s3_key)
48
+ const response = await this.s3.getObject(s3_input)
49
49
  const body = response.Body
50
50
  if (body === undefined) {
51
51
  return undefined
@@ -61,12 +61,12 @@ export class UtilsS3 {
61
61
  }
62
62
 
63
63
  async get_buffer(s3_id : string) : Promise<Buffer | undefined> {
64
- const s3_key = get_s3_key(s3_id)
65
- if (s3_key === undefined) {
64
+ const s3_input = get_s3_input(s3_id)
65
+ if (s3_input === undefined) {
66
66
  return undefined
67
67
  }
68
68
  try {
69
- const response = await this.s3.getObject(s3_key)
69
+ const response = await this.s3.getObject(s3_input)
70
70
  const body = response.Body
71
71
  if (body === undefined) {
72
72
  return undefined
@@ -82,12 +82,12 @@ export class UtilsS3 {
82
82
  }
83
83
 
84
84
  async get_stream(s3_id : string) : Promise<Readable | undefined> {
85
- const s3_key = get_s3_key(s3_id)
86
- if (s3_key === undefined) {
85
+ const s3_input = get_s3_input(s3_id)
86
+ if (s3_input === undefined) {
87
87
  return undefined
88
88
  }
89
89
  try {
90
- const response = await this.s3.getObject(s3_key)
90
+ const response = await this.s3.getObject(s3_input)
91
91
  const body = response.Body
92
92
  if (body === undefined) {
93
93
  return undefined
@@ -106,17 +106,17 @@ export class UtilsS3 {
106
106
 
107
107
  async query_prefix(s3_id_prefix : string, options : { compile? : boolean } = {}) {
108
108
  const compile = options.compile !== undefined ? options.compile : false
109
- const s3_key = get_s3_key(s3_id_prefix)
110
- if (s3_key === undefined) {
109
+ const s3_input = get_s3_input(s3_id_prefix)
110
+ if (s3_input === undefined) {
111
111
  return []
112
112
  }
113
113
  const s3_ids : string[] = []
114
114
  let last_token : string | undefined = undefined
115
115
  while (true) {
116
116
  const request : ListObjectsV2CommandInput = {
117
- Bucket : s3_key.Bucket,
117
+ Bucket : s3_input.Bucket,
118
118
  Delimiter : "/",
119
- Prefix : s3_key.Key,
119
+ Prefix : s3_input.Key,
120
120
  ContinuationToken : last_token
121
121
  }
122
122
  const response : ListObjectsV2CommandOutput = await this.s3.listObjectsV2(request)
@@ -124,12 +124,12 @@ export class UtilsS3 {
124
124
  .map(content => content.Key)
125
125
  .filter(key => key !== undefined)
126
126
  .filter(key => key.substring(key.length - 1) !== "/")
127
- .map(key => "s3://" + s3_key.Bucket + "/" + key)
127
+ .map(key => "s3://" + s3_input.Bucket + "/" + key)
128
128
  s3_ids.push(...new_s3_ids)
129
129
  const sub_s3_id_prefixes = (response.CommonPrefixes || [])
130
130
  .map(prefix => prefix.Prefix)
131
131
  .filter(prefix => prefix !== undefined)
132
- .map(prefix => "s3://" + s3_key.Bucket + "/" + prefix)
132
+ .map(prefix => "s3://" + s3_input.Bucket + "/" + prefix)
133
133
  for (const sub_s3_id_prefix of sub_s3_id_prefixes) {
134
134
  const new_sub_s3_ids = await this.query_prefix(sub_s3_id_prefix)
135
135
  s3_ids.push(...new_sub_s3_ids)
@@ -150,23 +150,23 @@ export class UtilsS3 {
150
150
  if (headers !== undefined) {
151
151
  return undefined
152
152
  }
153
- const s3_key = get_s3_key(s3_id)
154
- if (s3_key === undefined) {
153
+ const s3_input = get_s3_input(s3_id)
154
+ if (s3_input === undefined) {
155
155
  return undefined
156
156
  }
157
157
  return await this.s3.putObject({
158
- ...s3_key,
158
+ ...s3_input,
159
159
  Body : content,
160
160
  ContentType : content_type
161
161
  })
162
162
  }
163
163
 
164
164
  async delete(s3_id : string) : Promise<void>{
165
- const s3_key = get_s3_key(s3_id)
166
- if (s3_key === undefined) {
165
+ const s3_input = get_s3_input(s3_id)
166
+ if (s3_input === undefined) {
167
167
  return undefined
168
168
  }
169
- await this.s3.deleteObject(s3_key)
169
+ await this.s3.deleteObject(s3_input)
170
170
  }
171
171
 
172
172
  async get_download_url(s3_id : string) : Promise<string | undefined> {
@@ -174,21 +174,21 @@ export class UtilsS3 {
174
174
  if (headers === undefined) {
175
175
  return undefined
176
176
  }
177
- const s3_key = get_s3_key(s3_id)
178
- if (s3_key === undefined) {
177
+ const s3_input = get_s3_input(s3_id)
178
+ if (s3_input === undefined) {
179
179
  return undefined
180
180
  }
181
- const command = new GetObjectCommand(s3_key)
181
+ const command = new GetObjectCommand(s3_input)
182
182
  const url = await getSignedUrl(this.s3, command)
183
183
  return url
184
184
  }
185
185
 
186
186
  async get_upload_url(s3_id : string) : Promise<string | undefined> {
187
- const s3_key = get_s3_key(s3_id)
188
- if (s3_key === undefined) {
187
+ const s3_input = get_s3_input(s3_id)
188
+ if (s3_input === undefined) {
189
189
  return undefined
190
190
  }
191
- const command = new PutObjectCommand(s3_key)
191
+ const command = new PutObjectCommand(s3_input)
192
192
  const url = await getSignedUrl(this.s3, command)
193
193
  return url
194
194
  }
package/src/UtilsS3Vectors.ts CHANGED
@@ -1,24 +1,45 @@
1
1
  import { ListVectorsCommandInput, S3Vectors } from "@aws-sdk/client-s3vectors"
2
2
 
3
- interface VectorKey {
4
- vector_bucket_name : string,
5
- vector_index_name : string,
6
- vector_key : string
7
- }
3
+ export class S3Vector {
4
+ readonly s3_vector_id : string
5
+ readonly coordinates : number[]
6
+ readonly headers : Record<string, any>
7
+
8
+ constructor(s3_vector : S3Vector) {
9
+ if (!S3Vector.is(s3_vector)) {
10
+ throw Error("Invalid input.")
11
+ }
12
+ this.s3_vector_id = s3_vector.s3_vector_id
13
+ this.coordinates = s3_vector.coordinates
14
+ this.headers = s3_vector.headers
15
+ }
8
16
 
9
- function get_vector_key(vector_id : string) : VectorKey | undefined {
10
- const [vector_bucket_name, vector_index_name, vector_key] = (vector_id.match(/^s3vectors\:\/\/([^\/]+)\/([^\/]+)\/([^$]+)$/) || []).slice(1, 4)
11
- if (vector_bucket_name === undefined || vector_index_name === undefined || vector_key === undefined) {
12
- return undefined
17
+ static is(s3_vector : any) : s3_vector is S3Vector {
18
+ return (
19
+ s3_vector !== undefined &&
20
+ typeof s3_vector.s3_vector_id === "string" &&
21
+ Array.isArray(s3_vector.coordinates) && s3_vector.coordinates.every((coordinate : any) => typeof coordinate === "number") &&
22
+ typeof s3_vector.headers === "object" && Object.keys(s3_vector.headers).every((header_id : any) => typeof header_id === "string")
23
+ )
13
24
  }
14
- return { vector_bucket_name : vector_bucket_name, vector_index_name : vector_index_name, vector_key : vector_key }
15
25
  }
16
26
 
17
- export interface Vector {
18
-
19
- key : string,
20
- data : number[],
21
- metadata : any
27
+ interface S3VectorsInput {
28
+ s3_vector_bucket_name : string,
29
+ s3_vector_index_name : string,
30
+ s3_vector_key : string
31
+ }
32
+
33
+ function get_s3vectors_input(s3_vector_id : string) : S3VectorsInput {
34
+ const [s3_vector_bucket_name, s3_vector_index_name, s3_vector_key] = (s3_vector_id.match(/^s3vectors:\/\/([^\/]+)\/([^\/]+)\/([^$]*)$/) || []).slice(1, 4)
35
+ if (s3_vector_bucket_name === undefined || s3_vector_index_name === undefined || s3_vector_key === undefined) {
36
+ throw Error("Bad s3_vector_id: " + s3_vector_id)
37
+ }
38
+ return {
39
+ s3_vector_bucket_name : s3_vector_bucket_name,
40
+ s3_vector_index_name : s3_vector_index_name,
41
+ s3_vector_key : s3_vector_key
42
+ }
22
43
  }
23
44
 
24
45
  export class UtilsS3Vectors {
@@ -29,13 +50,14 @@ export class UtilsS3Vectors {
29
50
  this.s3vectors = new S3Vectors({ region : region })
30
51
  }
31
52
 
32
- async scan(vector_bucket : string, vector_index : string) : Promise<string[]> {
33
- const vector_keys = []
53
+ async scan(s3_vector_id_prefix : string) : Promise<string[]> {
54
+ const { s3_vector_bucket_name, s3_vector_index_name } = get_s3vectors_input(s3_vector_id_prefix)
55
+ const s3_vector_ids : string[] = []
34
56
  let next_token : string | undefined = undefined
35
57
  while (true) {
36
58
  const request : ListVectorsCommandInput = {
37
- vectorBucketName : vector_bucket,
38
- indexName : vector_index,
59
+ vectorBucketName : s3_vector_bucket_name,
60
+ indexName : s3_vector_index_name,
39
61
  maxResults : 1000,
40
62
  nextToken : next_token
41
63
  }
@@ -43,40 +65,58 @@ export class UtilsS3Vectors {
43
65
  if (response === undefined || response.vectors === undefined) {
44
66
  return []
45
67
  }
46
- vector_keys.push(...response.vectors.map(vector => vector.key))
68
+ s3_vector_ids.push(
69
+ ...response.vectors
70
+ .map(vector => vector.key)
71
+ .filter(s3_vector_key => s3_vector_key !== undefined)
72
+ .map(s3_vector_key => "s3vectors://" + s3_vector_bucket_name + "/" + s3_vector_index_name + "/" + s3_vector_key)
73
+ )
47
74
  next_token = response.nextToken
48
75
  if (next_token === undefined) {
49
- return vector_keys.filter(vector_key => vector_key !== undefined)
76
+ return s3_vector_ids
50
77
  }
51
78
  }
52
79
  }
53
80
 
54
- async get(vector_bucket : string, vector_index : string, vector_key : string) : Promise<Vector | undefined> {
81
+ async get(s3_vector_id : string) : Promise<S3Vector | undefined> {
82
+ const { s3_vector_bucket_name, s3_vector_index_name, s3_vector_key } = get_s3vectors_input(s3_vector_id)
55
83
  const response = await this.s3vectors.getVectors({
56
- vectorBucketName : vector_bucket,
57
- indexName : vector_index,
58
- keys : [vector_key],
84
+ vectorBucketName : s3_vector_bucket_name,
85
+ indexName : s3_vector_index_name,
86
+ keys : [s3_vector_key],
59
87
  returnData : true,
60
88
  returnMetadata : true
61
89
  })
62
90
  const vector = response.vectors?.[0]
63
- if (vector === undefined || vector.key === undefined || vector.data === undefined || vector.data.float32 === undefined || vector.metadata === undefined) {
91
+ if (vector === undefined) {
64
92
  return undefined
65
93
  }
66
- return {
67
- key : vector.key,
68
- data : vector.data.float32,
69
- metadata : vector.metadata
94
+ const coordinates : any = vector.data?.float32
95
+ const headers : any = vector.metadata
96
+ const s3_vector : S3Vector = {
97
+ s3_vector_id : s3_vector_id,
98
+ coordinates : coordinates,
99
+ headers : headers
70
100
  }
101
+ if (!S3Vector.is(s3_vector)) {
102
+ console.log("Bad S3Vector:", s3_vector)
103
+ return undefined
104
+ }
105
+ return s3_vector
71
106
  }
72
107
 
73
- async batch_get(vector_bucket : string, vector_index : string, vector_keys : string[]) {
74
- const vectors = []
75
- for (let i = 0; i < vector_keys.length; i += 100) {
108
+ async batch_get(s3_vector_ids : string[]) {
109
+ const s3vectors_inputs : S3VectorsInput[] = s3_vector_ids.map(get_s3vectors_input)
110
+ const { s3_vector_bucket_name, s3_vector_index_name } = s3vectors_inputs[0]
111
+ if (!s3vectors_inputs.every(s3vectors_input => s3vectors_input.s3_vector_bucket_name === s3_vector_bucket_name && s3vectors_input.s3_vector_index_name === s3_vector_index_name)) {
112
+ return []
113
+ }
114
+ const s3_vectors : (S3Vector | undefined)[] = []
115
+ for (let i = 0; i < s3vectors_inputs.length; i += 100) {
76
116
  const new_vectors = await this.s3vectors.getVectors({
77
- vectorBucketName : vector_bucket,
78
- indexName : vector_index,
79
- keys : vector_keys.slice(i, i + 100),
117
+ vectorBucketName : s3_vector_bucket_name,
118
+ indexName : s3_vector_index_name,
119
+ keys : s3vectors_inputs.slice(i, i + 100).map(s3vectors_input => s3vectors_input.s3_vector_key),
80
120
  returnData : true,
81
121
  returnMetadata : true
82
122
  })
@@ -84,83 +124,125 @@ export class UtilsS3Vectors {
84
124
  if (new_vectors === undefined) {
85
125
  return []
86
126
  }
87
- const vectors_by_key : Record<string, Vector | undefined> = Object.fromEntries(new_vectors
88
- .map(vector => [
89
- vector.key,
90
- vector.key !== undefined && vector.data !== undefined && vector.data.float32 !== undefined && vector.metadata !== undefined ? {
91
- key : vector.key,
92
- data : vector.data.float32,
93
- metadata : vector.metadata
94
- } : undefined
95
- ])
96
- )
97
- vectors.push(...vector_keys.map(vector_key => vectors_by_key[vector_key]))
127
+ const new_s3_vectors : S3Vector[] = []
128
+ for (const new_vector of new_vectors) {
129
+ const s3_vector_key : any = new_vector.key
130
+ const coordinates : any = new_vector.data?.float32
131
+ const headers : any = new_vector.metadata
132
+ const new_s3_vector : S3Vector | undefined = {
133
+ s3_vector_id : "s3vectors://" + s3_vector_bucket_name + "/" + s3_vector_index_name + "/" + s3_vector_key,
134
+ coordinates : coordinates,
135
+ headers : headers
136
+ }
137
+ if (!S3Vector.is(new_s3_vector)) {
138
+ console.log("Bad S3Vector:", new_s3_vector)
139
+ continue
140
+ }
141
+ new_s3_vectors.push(new_s3_vector)
142
+ }
143
+ const new_s3_vectors_by_s3_vector_id = Object.fromEntries(new_s3_vectors.map(new_s3_vector => [new_s3_vector.s3_vector_id, new_s3_vector]))
144
+ s3_vectors.push(...s3_vector_ids.slice(i, i + 100).map(s3_vector_id => new_s3_vectors_by_s3_vector_id[s3_vector_id]))
98
145
  }
99
- return vectors
146
+ return s3_vectors
100
147
  }
101
148
 
102
- async create(vector_bucket : string, vector_index : string, vector : Vector) : Promise<void> {
149
+ async put(s3_vector : S3Vector) : Promise<void> {
150
+ const { s3_vector_bucket_name, s3_vector_index_name, s3_vector_key } = get_s3vectors_input(s3_vector.s3_vector_id)
103
151
  await this.s3vectors.putVectors({
104
- vectorBucketName : vector_bucket,
105
- indexName : vector_index,
152
+ vectorBucketName : s3_vector_bucket_name,
153
+ indexName : s3_vector_index_name,
106
154
  vectors : [{
107
- key : vector.key,
155
+ key : s3_vector_key,
108
156
  data : {
109
- float32 : vector.data
157
+ float32 : s3_vector.coordinates
110
158
  },
111
- metadata : vector.metadata
159
+ metadata : s3_vector.headers
112
160
  }]
113
161
  })
114
162
  }
115
163
 
116
- async batch_create(vector_bucket : string, vector_index : string, vectors : Vector[]) : Promise<void> {
117
- for (let i = 0; i < vectors.length; i += 100) {
164
+ async batch_put(s3_vectors : S3Vector[]) : Promise<void> {
165
+ const s3vectors_inputs : S3VectorsInput[] = s3_vectors.map(s3_vector => get_s3vectors_input(s3_vector.s3_vector_id))
166
+ const { s3_vector_bucket_name, s3_vector_index_name } = s3vectors_inputs[0]
167
+ if (!s3vectors_inputs.every(s3vectors_input => s3vectors_input.s3_vector_bucket_name === s3_vector_bucket_name && s3vectors_input.s3_vector_index_name === s3_vector_index_name)) {
168
+ return undefined
169
+ }
170
+ for (let i = 0; i < s3_vectors.length; i += 100) {
118
171
  await this.s3vectors.putVectors({
119
- vectorBucketName : vector_bucket,
120
- indexName : vector_index,
121
- vectors : vectors.slice(i, i + 100).map(vector => ({
122
- key : vector.key,
172
+ vectorBucketName : s3_vector_bucket_name,
173
+ indexName : s3_vector_index_name,
174
+ vectors : s3_vectors.slice(i, i + 100).map(s3_vector => ({
175
+ key : get_s3vectors_input(s3_vector.s3_vector_id).s3_vector_key,
123
176
  data : {
124
- float32 : vector.data
177
+ float32 : s3_vector.coordinates
125
178
  },
126
- metadata : vector.metadata
179
+ metadata : s3_vector.headers
127
180
  }))
128
181
  })
129
182
 
130
183
  }
131
184
  }
132
185
 
133
- async delete(vector_bucket : string, vector_index : string, vector_key : string) : Promise<void> {
186
+ async delete(s3_vector_id : string) : Promise<void> {
187
+ const { s3_vector_bucket_name, s3_vector_index_name, s3_vector_key } = get_s3vectors_input(s3_vector_id)
134
188
  await this.s3vectors.deleteVectors({
135
- vectorBucketName : vector_bucket,
136
- indexName : vector_index,
137
- keys : [vector_key]
189
+ vectorBucketName : s3_vector_bucket_name,
190
+ indexName : s3_vector_index_name,
191
+ keys : [s3_vector_key]
138
192
  })
139
193
  }
140
194
 
141
- async batch_delete(vector_bucket : string, vector_index : string, vector_keys : string[]) : Promise<void> {
142
- for (let i = 0; i < vector_keys.length; i += 100) {
195
+ async batch_delete(s3_vector_ids : string[]) : Promise<void> {
196
+ const s3vectors_inputs : S3VectorsInput[] = s3_vector_ids.map(get_s3vectors_input)
197
+ const { s3_vector_bucket_name, s3_vector_index_name } = s3vectors_inputs[0]
198
+ if (!s3vectors_inputs.every(s3vectors_input => s3vectors_input.s3_vector_bucket_name === s3_vector_bucket_name && s3vectors_input.s3_vector_index_name === s3_vector_index_name)) {
199
+ return undefined
200
+ }
201
+ for (let i = 0; i < s3vectors_inputs.length; i += 100) {
143
202
  await this.s3vectors.deleteVectors({
144
- vectorBucketName : vector_bucket,
145
- indexName : vector_index,
146
- keys : vector_keys.slice(i, i + 100)
203
+ vectorBucketName : s3_vector_bucket_name,
204
+ indexName : s3_vector_index_name,
205
+ keys : s3vectors_inputs.slice(i, i + 100).map(s3vectors_input => s3vectors_input.s3_vector_key)
147
206
  })
148
207
  }
149
208
  }
150
209
 
151
- async query(vector_bucket : string, vector_index : string, data : number[], num_vectors : number, filters : Record<string, any> | undefined = undefined) {
152
- const response = await this.s3vectors.queryVectors({
153
- vectorBucketName : vector_bucket,
154
- indexName : vector_index,
210
+ async query(s3_vector_id_prefix : string, coordinates : number[], num_vectors : number, filters : Record<string, any> | undefined = undefined) {
211
+ const { s3_vector_bucket_name, s3_vector_index_name } = get_s3vectors_input(s3_vector_id_prefix)
212
+ const new_vectors = await this.s3vectors.queryVectors({
213
+ vectorBucketName : s3_vector_bucket_name,
214
+ indexName : s3_vector_index_name,
155
215
  queryVector : {
156
- float32 : data
216
+ float32 : coordinates
157
217
  },
158
- topK : num_vectors,
218
+ topK : num_vectors,
159
219
  returnDistance : true,
160
220
  returnMetadata : true,
161
221
  filter : filters
162
222
  })
163
- return response.vectors
223
+ .then(response => response.vectors)
224
+ if (new_vectors === undefined) {
225
+ return []
226
+ }
227
+ const results : { s3_vector_id : string, distance : number }[] = []
228
+ for (const new_vector of new_vectors) {
229
+ const s3_vector_key = new_vector.key
230
+ if (s3_vector_key === undefined) {
231
+ console.log("Bad s3_vector_key:", s3_vector_key)
232
+ continue
233
+ }
234
+ const s3_vector_id = s3_vector_id_prefix + "/" + s3_vector_key
235
+ const distance = new_vector.distance
236
+ if (distance === undefined) {
237
+ console.log("Bad distance:", distance)
238
+ continue
239
+ }
240
+ results.push({
241
+ s3_vector_id : s3_vector_id,
242
+ distance : distance
243
+ })
244
+ }
245
+ return results
164
246
  }
165
247
 
166
248
  }