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 {
13 createInternalLabels,
14 diffLabels,
15 hasAlchemyLabels,
16 stripInternalLabels,
17 toLabels,
18} from "../Labels.ts";
19import type { Providers } from "../Providers.ts";
20import {
21 compact,
22 normalizeLocation,
23 parentOf,
24 parseName,
25 stableJson,
26 toPhysicalId,
27 listLocations,
28} from "./names.ts";
29import { resourceNameFromOperation, waitForOperation } from "./operations.ts";
30
31export type ReasoningEngineSpec =
32 aiplatform.GoogleCloudAiplatformV1ReasoningEngineSpec;
33export type ReasoningEngineContextSpec =
34 aiplatform.GoogleCloudAiplatformV1ReasoningEngineContextSpec;
35export type ReasoningEngineEncryptionSpec =
36 aiplatform.GoogleCloudAiplatformV1EncryptionSpec;
37
38export type ReasoningEngineProps = {
39 /**
40 * Engine id (the `{reasoning_engine}` segment). Vertex assigns this on
41 * create. Pass the value from a previous deploy to target the same
42 * engine. Immutable — changing it replaces the engine.
43 */
44 reasoningEngineId?: string;
45 /**
46 * Vertex AI location. Immutable — changing it replaces the engine.
47 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
48 */
49 location?: string;
50 /**
51 * Display name. Required by Vertex. Defaults to a generated id.
52 */
53 displayName?: string;
54 /**
55 * Human-readable description.
56 */
57 description?: string;
58 /**
59 * User labels. Alchemy ownership labels are merged in automatically.
60 */
61 labels?: Record<string, string>;
62 /**
63 * Engine configurations (framework, source, package, container).
64 */
65 spec?: ReasoningEngineSpec;
66 /**
67 * Context management (Memory Bank) configuration.
68 */
69 contextSpec?: ReasoningEngineContextSpec;
70 /**
71 * Customer-managed encryption key. Immutable — changing it replaces
72 * the engine.
73 */
74 encryptionSpec?: ReasoningEngineEncryptionSpec;
75};
76
77export type ReasoningEngine = Resource<
78 "GCP.AIPlatform.ReasoningEngine",
79 ReasoningEngineProps,
80 {
81 /** Full resource name. */
82 name: string;
83 /** Engine id (last path segment). */
84 reasoningEngineId: string;
85 /** Location id. */
86 location: string;
87 /** Project id. */
88 project: string;
89 /** Display name. */
90 displayName: string | undefined;
91 /** Description. */
92 description: string | undefined;
93 /** User labels (Alchemy ownership labels stripped). */
94 labels: Record<string, string>;
95 /** Agent framework. */
96 agentFramework: string | undefined;
97 /** RFC3339 creation timestamp. */
98 createTime: string | undefined;
99 /** RFC3339 last-update timestamp. */
100 updateTime: string | undefined;
101 /** Etag for read-modify-write. */
102 etag: string | undefined;
103 },
104 never,
105 Providers
106>;
107
108/**
109 * A Vertex AI Reasoning Engine (Agent Engine) — a customizable runtime
110 * for models that choose which actions to take and in which order.
111 *
112 * Changing `reasoningEngineId`, `location`, or `encryptionSpec` replaces
113 * the engine. Display name, description, labels, and spec update in place.
114 *
115 * ### Creating a Reasoning Engine
116 * **Example:** Display name and labels
117 * ```typescript
118 * const engine = yield* GCP.AIPlatform.ReasoningEngine("Agent", {
119 * displayName: "support-agent",
120 * labels: { env: "prod" },
121 * spec: { agentFramework: "custom" },
122 * });
123 * ```
124 *
125 * ### Querying
126 * **Example:** Call the default query method
127 * ```typescript
128 * const query = yield* GCP.AIPlatform.QueryReasoningEngine(engine);
129 * const result = yield* query({ body: { input: { input: "hello" } } });
130 * ```
131 *
132 * @resource
133 * @category AIPlatform
134 */
135export const ReasoningEngine = Resource<ReasoningEngine>(
136 "GCP.AIPlatform.ReasoningEngine",
137);
138
139export class ReasoningEngineNotResolved extends Data.TaggedError(
140 "GCP.AIPlatform.ReasoningEngineNotResolved",
141)<{
142 name: string;
143}> {}
144
145export class ReasoningEngineStillExists extends Data.TaggedError(
146 "GCP.AIPlatform.ReasoningEngineStillExists",
147)<{
148 name: string;
149}> {}
150
151const resourceName = (
152 project: string,
153 location: string,
154 reasoningEngineId: string,
155) =>
156 `projects/${project}/locations/${location}/reasoningEngines/${reasoningEngineId}`;
157
158const userLabels = (
159 labels: Record<string, string | undefined> | null | undefined,
160): Record<string, string> => stripInternalLabels(tagRecord(labels));
161
162const toAttrs = (
163 engine: aiplatform.GoogleCloudAiplatformV1ReasoningEngine,
164 project: string,
165) => {
166 const name = engine.name ?? "";
167 const parsed = parseName(name, "reasoningEngines");
168 return {
169 name,
170 reasoningEngineId: parsed.resourceId,
171 location: parsed.location,
172 project: parsed.project || project,
173 displayName: engine.displayName,
174 description: engine.description,
175 labels: userLabels(engine.labels),
176 agentFramework: engine.spec?.agentFramework,
177 createTime: engine.createTime,
178 updateTime: engine.updateTime,
179 etag: engine.etag,
180 };
181};
182
183const getByName = (name: string) =>
184 name.length === 0
185 ? Effect.succeed(undefined)
186 : aiplatform
187 .getProjectsLocationsReasoningEngines({ name })
188 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
189
190const listAt = (parent: string) =>
191 aiplatform.listProjectsLocationsReasoningEngines
192 .pages({ parent, pageSize: 100 })
193 .pipe(
194 Stream.flatMap((page) =>
195 Stream.fromIterable(page.reasoningEngines ?? []),
196 ),
197 Stream.take(500),
198 Stream.runCollect,
199 Effect.map((chunk) => Array.from(chunk)),
200 Effect.catchTag("NotFound", () => Effect.succeed([])),
201 );
202
203export const listAlchemyReasoningEngines = (
204 project: string,
205 location: string,
206) =>
207 listAt(parentOf(project, location)).pipe(
208 Effect.map((engines) =>
209 engines.filter((engine) =>
210 Object.keys(engine.labels ?? {}).some((key) =>
211 key.startsWith("alchemy-"),
212 ),
213 ),
214 ),
215 );
216
217const findOwned = (id: string, parent: string) =>
218 Effect.gen(function* () {
219 const expected = yield* createInternalLabels(id);
220 const engines = yield* listAt(parent);
221 return engines.find((engine) =>
222 Object.entries(expected).every(
223 ([key, value]) => (engine.labels ?? {})[key] === value,
224 ),
225 );
226 });
227
228const waitUntilGone = (name: string) =>
229 getByName(name).pipe(
230 Effect.flatMap((engine) =>
231 engine === undefined
232 ? Effect.void
233 : Effect.fail(new ReasoningEngineStillExists({ name })),
234 ),
235 Effect.retry({
236 while: (error) =>
237 error._tag === "GCP.AIPlatform.ReasoningEngineStillExists",
238 times: 10,
239 schedule: Schedule.spaced("2 seconds"),
240 }),
241 );
242
243// Engine writes share a small per-minute regional quota ("Reasoning Engine
244// Write Requests per minute per region"); wait out the window.
245const retryWriteQuota = <A, E extends { readonly _tag: string }, R>(
246 effect: Effect.Effect<A, E, R>,
247) =>
248 effect.pipe(
249 Effect.retry({
250 while: (error) => error._tag === "TooManyRequests",
251 times: 8,
252 schedule: Schedule.spaced("15 seconds"),
253 }),
254 );
255
256export const ReasoningEngineProvider = () =>
257 Provider.succeed(ReasoningEngine, {
258 stables: ["name", "reasoningEngineId", "location", "project", "createTime"],
259
260 diff: Effect.fn(function* ({ news, olds, output }) {
261 const env = yield* GcpEnvironment.current;
262 if (!isResolved(news)) return undefined;
263 const previousId = olds?.reasoningEngineId ?? output?.reasoningEngineId;
264 const nextId = news.reasoningEngineId ?? previousId;
265 const idChanged =
266 previousId !== undefined &&
267 nextId !== undefined &&
268 nextId !== previousId;
269 const previousLocation = normalizeLocation(
270 olds?.location ?? output?.location,
271 env.region,
272 );
273 const nextLocation = normalizeLocation(
274 news.location ?? output?.location,
275 env.region,
276 );
277 const encryptionChanged =
278 (news.encryptionSpec?.kmsKeyName ?? "") !==
279 (olds?.encryptionSpec?.kmsKeyName ?? "");
280 if (
281 idChanged ||
282 previousLocation !== nextLocation ||
283 (olds !== undefined && encryptionChanged)
284 ) {
285 return { action: "replace" as const, deleteFirst: false };
286 }
287 return undefined;
288 }),
289
290 read: Effect.fn(function* ({ id, olds, output }) {
291 const env = yield* GcpEnvironment.current;
292 const location = normalizeLocation(
293 olds?.location ?? output?.location,
294 env.region,
295 );
296 const name =
297 output?.name ??
298 (olds?.reasoningEngineId !== undefined
299 ? resourceName(env.project, location, olds.reasoningEngineId)
300 : "");
301 const existing =
302 (name.length > 0 ? yield* getByName(name) : undefined) ??
303 (yield* findOwned(id, parentOf(env.project, location)));
304 if (existing === undefined) return undefined;
305 const attrs = toAttrs(existing, env.project);
306 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
307 ? attrs
308 : Unowned(attrs);
309 }),
310
311 list: () =>
312 Effect.gen(function* () {
313 const env = yield* GcpEnvironment.current;
314 const engines = (yield* Effect.forEach(
315 listLocations(env.region),
316 (location) => listAt(parentOf(env.project, location)),
317 )).flat();
318 return engines
319 .filter((engine) =>
320 Object.keys(engine.labels ?? {}).some((key) =>
321 key.startsWith("alchemy-"),
322 ),
323 )
324 .map((engine) => toAttrs(engine, env.project));
325 }),
326
327 reconcile: Effect.fn(function* ({ id, news, output }) {
328 const env = yield* GcpEnvironment.current;
329 const location = normalizeLocation(
330 news.location ?? output?.location,
331 env.region,
332 );
333 const parent = parentOf(env.project, location);
334 const desiredLabels = {
335 ...toLabels(news.labels),
336 ...(yield* createInternalLabels(id)),
337 };
338 const displayName =
339 news.displayName ??
340 (yield* toPhysicalId(
341 id,
342 news.reasoningEngineId,
343 output?.reasoningEngineId,
344 ));
345
346 let current =
347 (output?.name !== undefined
348 ? yield* getByName(output.name)
349 : undefined) ?? (yield* findOwned(id, parent));
350
351 if (current === undefined) {
352 const created = yield* retryWriteQuota(
353 aiplatform.createProjectsLocationsReasoningEngines({
354 parent,
355 body: compact({
356 displayName,
357 description: news.description,
358 labels: desiredLabels,
359 spec: news.spec,
360 contextSpec: news.contextSpec,
361 encryptionSpec: news.encryptionSpec,
362 }),
363 }),
364 ).pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
365 if (created !== undefined) {
366 yield* waitForOperation(created, { alreadyExistsOk: true });
367 }
368 const createdName = resourceNameFromOperation(created ?? {});
369 current =
370 createdName !== undefined
371 ? yield* getByName(createdName)
372 : yield* findOwned(id, parent);
373 }
374
375 if (current === undefined || current.name === undefined) {
376 return yield* new ReasoningEngineNotResolved({
377 name: output?.name ?? parent,
378 });
379 }
380
381 const observedName = current.name;
382 const observedLabels = tagRecord(current.labels);
383 const { upsert, removed } = diffLabels(observedLabels, desiredLabels);
384 const labelsChanged = upsert.length > 0 || removed.length > 0;
385 const displayChanged = (current.displayName ?? "") !== displayName;
386 const descriptionChanged =
387 (current.description ?? "") !== (news.description ?? "");
388 const specChanged =
389 stableJson(current.spec) !== stableJson(news.spec ?? current.spec);
390
391 if (
392 labelsChanged ||
393 displayChanged ||
394 descriptionChanged ||
395 specChanged
396 ) {
397 const updateMask = [
398 labelsChanged ? "labels" : undefined,
399 displayChanged ? "display_name" : undefined,
400 descriptionChanged ? "description" : undefined,
401 specChanged ? "spec" : undefined,
402 ].filter((field): field is string => field !== undefined);
403 const patched = yield* retryWriteQuota(
404 aiplatform.patchProjectsLocationsReasoningEngines({
405 name: observedName,
406 updateMask: updateMask.join(","),
407 body: compact({
408 name: observedName,
409 displayName,
410 description: news.description,
411 labels: desiredLabels,
412 spec: news.spec,
413 etag: current.etag,
414 }),
415 }),
416 );
417 yield* waitForOperation(patched);
418 current = yield* getByName(observedName);
419 }
420
421 if (current === undefined) {
422 return yield* new ReasoningEngineNotResolved({ name: observedName });
423 }
424 return toAttrs(current, env.project);
425 }),
426
427 delete: Effect.fn(function* ({ output }) {
428 const operation = yield* retryWriteQuota(
429 aiplatform.deleteProjectsLocationsReasoningEngines({
430 name: output.name,
431 force: true,
432 }),
433 ).pipe(
434 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
435 Effect.catchTag("BadRequest", () => Effect.succeed(undefined)),
436 );
437 if (operation !== undefined) {
438 yield* waitForOperation(operation, { notFoundOk: true });
439 }
440 yield* waitUntilGone(output.name);
441 }),
442 });