1import * as aiplatform from "@distilled.cloud/gcp/aiplatform_v1";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
4import * as Schedule from "effect/Schedule";
5import * as Stream from "effect/Stream";
6import { Unowned } from "../../AdoptPolicy.ts";
7import { isResolved } from "../../Diff.ts";
8import * as Provider from "../../Provider.ts";
9import { Resource } from "../../Resource.ts";
10import { tagRecord } from "../../Tags.ts";
11import { GcpEnvironment } from "../Environment.ts";
12import { hasAlchemyLabels } from "../Labels.ts";
13import type { Providers } from "../Providers.ts";
14import {
15 alchemyIdFilter,
16 createInternalLabels,
17 hasAlchemyPrefix,
18 labelsDiffer,
19 listLocations,
20 locationOf,
21 lastSegment,
22 normalizeLocation,
23 projectOf,
24 resourceNameFromOperation,
25 stableJson,
26 toDisplayName,
27 toLabels,
28 userLabels,
29} from "./names.ts";
30import { waitForOperation } from "./operations.ts";
31
32export type IndexProps = {
33 /**
34 * Vertex AI location (`us-central1`, …). Immutable — changing it
35 * replaces the index.
36 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
37 */
38 location?: string;
39 /**
40 * Display name (max 128 Unicode characters). Generated from the stack,
41 * stage, and logical id when omitted.
42 */
43 displayName?: string;
44 /**
45 * Human-readable description.
46 */
47 description?: string;
48 /**
49 * User labels. Alchemy ownership labels are merged in automatically.
50 */
51 labels?: Record<string, string>;
52 /**
53 * How the index is updated. Immutable.
54 * @default "STREAM_UPDATE"
55 */
56 indexUpdateMethod?:
57 | aiplatform.GoogleCloudAiplatformV1IndexIndexUpdateMethodEnum
58 | (string & {});
59 /**
60 * Matching Engine metadata (`contentsDeltaUri`, `config.dimensions`,
61 * algorithm config).
62 */
63 metadata?: unknown;
64 /**
65 * GCS URI of the OpenAPI schema for `metadata`. Immutable.
66 */
67 metadataSchemaUri?: string;
68 /**
69 * Customer-managed encryption key. Immutable.
70 */
71 encryptionSpec?: aiplatform.GoogleCloudAiplatformV1EncryptionSpec;
72};
73
74export type Index = Resource<
75 "GCP.AIPlatform.Index",
76 IndexProps,
77 {
78 /** Full resource name `projects/{project}/locations/{location}/indexes/{index}`. */
79 name: string;
80 /** Index id (last path segment). */
81 indexId: string;
82 /** Project id. */
83 project: string;
84 /** Location id. */
85 location: string;
86 /** Display name. */
87 displayName: string | undefined;
88 /** User-provided description. */
89 description: string | undefined;
90 /** User labels (Alchemy ownership labels stripped). */
91 labels: Record<string, string>;
92 /** Update method (`BATCH_UPDATE` or `STREAM_UPDATE`). */
93 indexUpdateMethod: string | undefined;
94 /** Matching Engine metadata. */
95 metadata: unknown;
96 /** Schema URI for metadata, if set. */
97 metadataSchemaUri: string | undefined;
98 /** Dense vector count. */
99 vectorsCount: string | undefined;
100 /** RFC3339 creation timestamp. */
101 createTime: string | undefined;
102 /** RFC3339 last-update timestamp. */
103 updateTime: string | undefined;
104 },
105 never,
106 Providers
107>;
108
109/**
110 * A Vertex AI Matching Engine Index of embedding vectors.
111 *
112 * Changing `location`, `indexUpdateMethod`, encryption, or
113 * `metadataSchemaUri` replaces the index. Display name, description,
114 * labels, and metadata update in place (metadata updates are long-running).
115 *
116 * ### Creating an Index
117 * **Example:** Empty streaming index
118 * ```typescript
119 * const index = yield* GCP.AIPlatform.Index("Embeddings", {
120 * displayName: "product-embeddings",
121 * indexUpdateMethod: "STREAM_UPDATE",
122 * metadata: {
123 * config: {
124 * dimensions: 768,
125 * approximateNeighborsCount: 150,
126 * distanceMeasureType: "DOT_PRODUCT_DISTANCE",
127 * shardSize: "SHARD_SIZE_SMALL",
128 * algorithmConfig: { bruteForceConfig: {} },
129 * },
130 * },
131 * labels: { env: "prod" },
132 * });
133 * ```
134 *
135 * @resource
136 * @category AIPlatform
137 */
138export const Index = Resource<Index>("GCP.AIPlatform.Index");
139
140export class IndexNotResolved extends Data.TaggedError(
141 "GCP.AIPlatform.IndexNotResolved",
142)<{
143 name: string;
144}> {}
145
146export class IndexStillExists extends Data.TaggedError(
147 "GCP.AIPlatform.IndexStillExists",
148)<{
149 name: string;
150}> {}
151
152const DEFAULT_UPDATE_METHOD = "STREAM_UPDATE";
153
154const toAttrs = (
155 index: aiplatform.GoogleCloudAiplatformV1Index,
156 project: string,
157) => {
158 const name = index.name ?? "";
159 return {
160 name,
161 indexId: lastSegment(name),
162 project: projectOf(name, project),
163 location: locationOf(name),
164 displayName: index.displayName,
165 description: index.description,
166 labels: userLabels(index.labels),
167 indexUpdateMethod: index.indexUpdateMethod,
168 metadata: index.metadata,
169 metadataSchemaUri: index.metadataSchemaUri,
170 vectorsCount: index.indexStats?.vectorsCount,
171 createTime: index.createTime,
172 updateTime: index.updateTime,
173 };
174};
175
176const getByName = (name: string) =>
177 aiplatform
178 .getProjectsLocationsIndexes({ name })
179 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
180
181const listPage = (parent: string, filter?: string) =>
182 aiplatform.listProjectsLocationsIndexes
183 .pages({ parent, pageSize: 100, filter })
184 .pipe(
185 Stream.runCollect,
186 Effect.map((pages) =>
187 Array.from(pages).flatMap((page) => page.indexes ?? []),
188 ),
189 Effect.catchTag("NotFound", () =>
190 Effect.succeed([] as aiplatform.GoogleCloudAiplatformV1Index[]),
191 ),
192 );
193
194const findOwned = (
195 project: string,
196 location: string,
197 labels: Record<string, string>,
198) =>
199 listPage(
200 `projects/${project}/locations/${location}`,
201 alchemyIdFilter(labels),
202 ).pipe(
203 Effect.map(
204 (items) =>
205 items.find((item) => hasAlchemyPrefix(item.labels)) ?? undefined,
206 ),
207 );
208
209const waitUntilExists = (name: string) =>
210 getByName(name).pipe(
211 Effect.filterOrFail(
212 (index): index is aiplatform.GoogleCloudAiplatformV1Index =>
213 index !== undefined,
214 () => new IndexNotResolved({ name }),
215 ),
216 Effect.retry({
217 while: (error) => error._tag === "GCP.AIPlatform.IndexNotResolved",
218 times: 8,
219 schedule: Schedule.spaced("2 seconds"),
220 }),
221 );
222
223const waitUntilGone = (name: string) =>
224 getByName(name).pipe(
225 Effect.filterOrFail(
226 (index) => index === undefined,
227 () => new IndexStillExists({ name }),
228 ),
229 Effect.asVoid,
230 Effect.retry({
231 while: (error) => error._tag === "GCP.AIPlatform.IndexStillExists",
232 times: 10,
233 schedule: Schedule.spaced("2 seconds"),
234 }),
235 );
236
237export const IndexProvider = () =>
238 Provider.succeed(Index, {
239 stables: [
240 "name",
241 "indexId",
242 "project",
243 "location",
244 "indexUpdateMethod",
245 "createTime",
246 ],
247
248 diff: Effect.fn(function* ({ news, olds, output }) {
249 const env = yield* GcpEnvironment.current;
250 if (!isResolved(news)) return undefined;
251 const previousLocation = normalizeLocation(
252 olds?.location ?? output?.location,
253 env.region,
254 );
255 const nextLocation = normalizeLocation(
256 news.location ?? olds?.location ?? output?.location,
257 env.region,
258 );
259 const previousMethod = (
260 olds?.indexUpdateMethod ??
261 output?.indexUpdateMethod ??
262 DEFAULT_UPDATE_METHOD
263 ).toUpperCase();
264 const nextMethod = (
265 news.indexUpdateMethod ?? previousMethod
266 ).toUpperCase();
267 const previousSchema =
268 olds?.metadataSchemaUri ?? output?.metadataSchemaUri ?? "";
269 const nextSchema = news.metadataSchemaUri ?? previousSchema;
270 if (
271 previousLocation !== nextLocation ||
272 previousMethod !== nextMethod ||
273 nextSchema !== previousSchema
274 ) {
275 return { action: "replace" as const, deleteFirst: false };
276 }
277 return undefined;
278 }),
279
280 read: Effect.fn(function* ({ id, olds, output }) {
281 const env = yield* GcpEnvironment.current;
282 const location = normalizeLocation(
283 olds?.location ?? output?.location,
284 env.region,
285 );
286 const ownership = yield* createInternalLabels(id);
287 const existing =
288 (output?.name !== undefined
289 ? yield* getByName(output.name)
290 : undefined) ?? (yield* findOwned(env.project, location, ownership));
291 if (existing === undefined) return undefined;
292 const attrs = toAttrs(existing, env.project);
293 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
294 ? attrs
295 : Unowned(attrs);
296 }),
297
298 list: () =>
299 Effect.gen(function* () {
300 const env = yield* GcpEnvironment.current;
301 const pages = yield* Effect.forEach(
302 listLocations(env.region),
303 (location) =>
304 listPage(`projects/${env.project}/locations/${location}`),
305 { concurrency: 4 },
306 );
307 return pages
308 .flat()
309 .filter((index) => hasAlchemyPrefix(index.labels))
310 .map((index) => toAttrs(index, env.project));
311 }),
312
313 reconcile: Effect.fn(function* ({ id, news, output }) {
314 const env = yield* GcpEnvironment.current;
315 const location = normalizeLocation(
316 news.location ?? output?.location,
317 env.region,
318 );
319 const displayName = yield* toDisplayName(
320 id,
321 news.displayName,
322 output?.displayName,
323 );
324 const desiredLabels = {
325 ...toLabels(news.labels),
326 ...(yield* createInternalLabels(id)),
327 };
328 const indexUpdateMethod = (
329 news.indexUpdateMethod ??
330 output?.indexUpdateMethod ??
331 DEFAULT_UPDATE_METHOD
332 ).toUpperCase();
333
334 let current =
335 (output?.name !== undefined
336 ? yield* getByName(output.name)
337 : undefined) ??
338 (yield* findOwned(env.project, location, desiredLabels));
339
340 if (current === undefined) {
341 const created = yield* aiplatform
342 .createProjectsLocationsIndexes({
343 parent: `projects/${env.project}/locations/${location}`,
344 body: {
345 displayName,
346 description: news.description,
347 labels: desiredLabels,
348 indexUpdateMethod,
349 metadata: news.metadata,
350 metadataSchemaUri: news.metadataSchemaUri,
351 encryptionSpec: news.encryptionSpec,
352 },
353 })
354 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
355 if (created !== undefined) {
356 const done = yield* waitForOperation(created);
357 const name = resourceNameFromOperation(done);
358 current =
359 name !== undefined
360 ? yield* waitUntilExists(name)
361 : yield* findOwned(env.project, location, desiredLabels);
362 }
363 if (current === undefined) {
364 current = yield* findOwned(env.project, location, desiredLabels);
365 }
366 }
367
368 if (current === undefined || current.name === undefined) {
369 return yield* new IndexNotResolved({
370 name: output?.name ?? `${location}/indexes`,
371 });
372 }
373
374 const name = current.name;
375 const descriptionChanged =
376 (current.description ?? "") !== (news.description ?? "");
377 const displayChanged = (current.displayName ?? "") !== displayName;
378 const labelsChanged = labelsDiffer(current.labels, desiredLabels);
379 const metadataChanged =
380 news.metadata !== undefined &&
381 stableJson(current.metadata) !== stableJson(news.metadata);
382
383 if (
384 descriptionChanged ||
385 displayChanged ||
386 labelsChanged ||
387 metadataChanged
388 ) {
389 const operation = yield* aiplatform.patchProjectsLocationsIndexes({
390 name,
391 updateMask: [
392 descriptionChanged ? "description" : undefined,
393 displayChanged ? "display_name" : undefined,
394 labelsChanged ? "labels" : undefined,
395 metadataChanged ? "metadata" : undefined,
396 ]
397 .filter((field): field is string => field !== undefined)
398 .join(","),
399 body: {
400 name,
401 displayName,
402 description: news.description,
403 labels: desiredLabels,
404 metadata: news.metadata,
405 },
406 });
407 yield* waitForOperation(operation);
408 current = yield* waitUntilExists(name);
409 }
410
411 return toAttrs(current, env.project);
412 }),
413
414 delete: Effect.fn(function* ({ output }) {
415 const operation = yield* aiplatform
416 .deleteProjectsLocationsIndexes({ name: output.name })
417 .pipe(
418 Effect.retry({
419 while: (error) => error._tag === "Conflict",
420 times: 8,
421 schedule: Schedule.spaced("2 seconds"),
422 }),
423 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
424 );
425 if (operation !== undefined) {
426 yield* waitForOperation(operation, { notFoundOk: true });
427 }
428 yield* waitUntilGone(output.name);
429 }),
430 });