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 { createPhysicalName } from "../../PhysicalName.ts";
9import * as Provider from "../../Provider.ts";
10import { Resource } from "../../Resource.ts";
11import { GcpEnvironment } from "../Environment.ts";
12import { createInternalLabels, hasAlchemyLabels } from "../Labels.ts";
13import type { Providers } from "../Providers.ts";
14import { listLocations } from "./names.ts";
15import {
16 encodeOwnership,
17 hasOwnershipMarker,
18 normalizeLocation,
19 parentOf,
20 parseOwnership,
21 parseResourceName,
22 rfc1035,
23} from "./internal.ts";
24import type { EncryptionSpec } from "./shared.ts";
25
26export type CachedContentPart = {
27 /** Text content of the part. */
28 text?: string;
29};
30
31export type CachedContentMessage = {
32 /**
33 * Producer of the content (`user` or `model`). Defaults to `user`.
34 */
35 role?: string;
36 /** Parts that make up the message. */
37 parts?: CachedContentPart[];
38};
39
40export type CachedContentProps = {
41 /**
42 * Region. Immutable — changing it replaces the cache.
43 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
44 */
45 location?: string;
46 /**
47 * Publisher model used for the cache, for example
48 * `projects/{project}/locations/{location}/publishers/google/models/gemini-2.0-flash-001`.
49 * Immutable — changing it replaces the cache.
50 */
51 model: string;
52 /**
53 * Content to cache. Immutable — changing it replaces the cache.
54 */
55 contents?: CachedContentMessage[];
56 /**
57 * Developer system instruction (text only). Immutable — changing it
58 * replaces the cache.
59 */
60 systemInstruction?: CachedContentMessage;
61 /**
62 * TTL such as `"3600s"`. The API computes `expireTime = now + ttl`.
63 * This is the only field that updates in place (along with
64 * `expireTime`).
65 */
66 ttl?: string;
67 /**
68 * Absolute expiration (RFC3339). Mutually exclusive with `ttl` on
69 * update; either field may be patched.
70 */
71 expireTime?: string;
72 /**
73 * User-facing display name. Alchemy stamps ownership here (cached
74 * content has no labels) so `list` / nuke can find the cache.
75 */
76 displayName?: string;
77 /**
78 * Customer-managed encryption. Immutable — changing it replaces the
79 * cache.
80 */
81 encryptionSpec?: EncryptionSpec;
82};
83
84export type CachedContent = Resource<
85 "GCP.AIPlatform.CachedContent",
86 CachedContentProps,
87 {
88 /** Full resource name `.../cachedContents/{cached_content}`. */
89 name: string;
90 /** Cached content id (last path segment). */
91 cachedContentId: string;
92 /** Project id. */
93 project: string;
94 /** Location id. */
95 location: string;
96 /** Model resource name. */
97 model: string | undefined;
98 /** User display name with the Alchemy ownership prefix stripped. */
99 displayName: string | undefined;
100 /** Absolute expiration (RFC3339). */
101 expireTime: string | undefined;
102 /** RFC3339 creation timestamp. */
103 createTime: string | undefined;
104 /** RFC3339 last-update timestamp. */
105 updateTime: string | undefined;
106 /** Total tokens consumed by the cache. */
107 totalTokenCount: number | undefined;
108 /** Customer-managed KMS key, if any. */
109 kmsKeyName: string | undefined;
110 },
111 never,
112 Providers
113>;
114
115/**
116 * Vertex AI cached content — explicit context cache for LLM queries.
117 *
118 * Cached content has no labels field, so Alchemy stamps ownership into
119 * `displayName` for `list` / nuke. Model, contents, system instruction,
120 * and CMEK are immutable. TTL / expire time update in place.
121 *
122 * ### Creating Cached Content
123 * **Example:** Cache a system prompt for Gemini
124 * ```typescript
125 * const cache = yield* GCP.AIPlatform.CachedContent("Style", {
126 * model:
127 * "projects/my-project/locations/us-central1/publishers/google/models/gemini-2.0-flash-001",
128 * ttl: "3600s",
129 * contents: [
130 * { role: "user", parts: [{ text: "You are a terse assistant." }] },
131 * ],
132 * });
133 * ```
134 *
135 * ### Updating expiry
136 * **Example:** Extend TTL
137 * ```typescript
138 * // Same logical id, changed TTL: the engine updates it in place.
139 * const cache = yield* GCP.AIPlatform.CachedContent("Style", {
140 * model:
141 * "projects/my-project/locations/us-central1/publishers/google/models/gemini-2.0-flash-001",
142 * ttl: "7200s",
143 * contents: [
144 * { role: "user", parts: [{ text: "You are a terse assistant." }] },
145 * ],
146 * });
147 * ```
148 *
149 * @resource
150 * @category AIPlatform
151 */
152export const CachedContent = Resource<CachedContent>(
153 "GCP.AIPlatform.CachedContent",
154);
155
156export class CachedContentNotResolved extends Data.TaggedError(
157 "GCP.AIPlatform.CachedContentNotResolved",
158)<{
159 name: string;
160}> {}
161
162const toContents = (
163 contents: readonly CachedContentMessage[] | undefined,
164): aiplatform.GoogleCloudAiplatformV1Content[] | undefined => {
165 if (contents === undefined) return undefined;
166 return contents.map((content) => ({
167 role: content.role,
168 parts: (content.parts ?? [])
169 .filter((part) => part.text !== undefined)
170 .map((part) => ({ text: part.text })),
171 }));
172};
173
174const toAttrs = (
175 cache: aiplatform.GoogleCloudAiplatformV1CachedContent,
176 project: string,
177) => {
178 const name = cache.name ?? "";
179 const parsed = parseResourceName(name, "cachedContents");
180 const ownership = parseOwnership(cache.displayName);
181 return {
182 name,
183 cachedContentId: parsed.id,
184 project: parsed.project || project,
185 location: parsed.location,
186 model: cache.model,
187 displayName: ownership.text,
188 expireTime: cache.expireTime,
189 createTime: cache.createTime,
190 updateTime: cache.updateTime,
191 totalTokenCount: cache.usageMetadata?.totalTokenCount,
192 kmsKeyName: cache.encryptionSpec?.kmsKeyName,
193 };
194};
195
196const getByName = (name: string) =>
197 aiplatform
198 .getProjectsLocationsCachedContents({ name })
199 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
200
201const listCaches = (project: string, region: string) => {
202 const collect = (parent: string) =>
203 aiplatform.listProjectsLocationsCachedContents
204 .pages({ parent, pageSize: 100 })
205 .pipe(
206 Stream.flatMap((page) =>
207 Stream.fromIterable(page.cachedContents ?? []),
208 ),
209 Stream.runCollect,
210 Effect.map((chunk) => Array.from(chunk)),
211 );
212 const fallback = Effect.forEach(listLocations(region), (location) =>
213 collect(`projects/${project}/locations/${location}`),
214 ).pipe(Effect.map((pages) => pages.flat()));
215 // Vertex AI has no `locations/-` wildcard; scan known locations.
216 return fallback.pipe(Effect.catchTag("NotFound", () => Effect.succeed([])));
217};
218
219const findOwned = (
220 id: string,
221 project: string,
222 region: string,
223 hinted?: string,
224) =>
225 Effect.gen(function* () {
226 if (hinted !== undefined && hinted.length > 0) {
227 const existing = yield* getByName(hinted);
228 if (existing !== undefined) return existing;
229 }
230 const caches = yield* listCaches(project, region);
231 for (const cache of caches) {
232 const { labels } = parseOwnership(cache.displayName);
233 if (yield* hasAlchemyLabels(id, labels)) return cache;
234 }
235 return undefined as
236 | aiplatform.GoogleCloudAiplatformV1CachedContent
237 | undefined;
238 });
239
240const fallbackDisplayName = (id: string) =>
241 Effect.gen(function* () {
242 return rfc1035(
243 yield* createPhysicalName({
244 id,
245 maxLength: 63,
246 lowercase: true,
247 }),
248 );
249 });
250
251export const CachedContentProvider = () =>
252 Provider.succeed(CachedContent, {
253 stables: [
254 "name",
255 "cachedContentId",
256 "project",
257 "location",
258 "model",
259 "createTime",
260 ],
261
262 diff: Effect.fn(function* ({ news, olds, output }) {
263 const env = yield* GcpEnvironment.current;
264 if (!isResolved(news)) return undefined;
265 const previousLocation = normalizeLocation(
266 olds?.location ?? output?.location,
267 env.region,
268 );
269 const nextLocation = normalizeLocation(
270 news.location ?? olds?.location ?? output?.location,
271 env.region,
272 );
273 const previousModel = olds?.model ?? output?.model ?? "";
274 const previousKey =
275 olds?.encryptionSpec?.kmsKeyName ?? output?.kmsKeyName ?? "";
276 const nextKey = news.encryptionSpec?.kmsKeyName ?? previousKey;
277 const replace =
278 previousLocation !== nextLocation ||
279 news.model !== previousModel ||
280 nextKey !== previousKey;
281 if (!replace) return undefined;
282 return { action: "replace" as const, deleteFirst: true };
283 }),
284
285 read: Effect.fn(function* ({ id, olds, output }) {
286 const env = yield* GcpEnvironment.current;
287 const existing = yield* findOwned(
288 id,
289 env.project,
290 env.region,
291 output?.name,
292 );
293 if (existing === undefined) return undefined;
294 const attrs = toAttrs(existing, env.project);
295 const { labels } = parseOwnership(existing.displayName);
296 return (yield* hasAlchemyLabels(id, labels)) ? attrs : Unowned(attrs);
297 }),
298
299 list: () =>
300 Effect.gen(function* () {
301 const env = yield* GcpEnvironment.current;
302 const caches = yield* listCaches(env.project, env.region);
303 return caches
304 .filter((cache) => hasOwnershipMarker(cache.displayName))
305 .map((cache) => toAttrs(cache, env.project));
306 }),
307
308 reconcile: Effect.fn(function* ({ id, news, output }) {
309 const env = yield* GcpEnvironment.current;
310 const location = normalizeLocation(
311 news.location ?? output?.location,
312 env.region,
313 );
314 const internal = yield* createInternalLabels(id);
315 const desiredDisplayName = encodeOwnership(
316 internal,
317 news.displayName ?? (yield* fallbackDisplayName(id)),
318 " ",
319 );
320
321 let current = yield* findOwned(id, env.project, env.region, output?.name);
322
323 if (current === undefined) {
324 const created = yield* aiplatform
325 .createProjectsLocationsCachedContents({
326 parent: parentOf(env.project, location),
327 body: {
328 model: news.model,
329 contents: toContents(news.contents),
330 systemInstruction: news.systemInstruction
331 ? toContents([news.systemInstruction])?.[0]
332 : undefined,
333 ttl: news.ttl,
334 expireTime: news.expireTime,
335 displayName: desiredDisplayName,
336 encryptionSpec: news.encryptionSpec,
337 },
338 })
339 .pipe(
340 // Vertex AI intermittently counts large content as "1 tokens".
341 Effect.retry({
342 while: (error) => error._tag === "CachedContentTooFewTokens",
343 times: 4,
344 schedule: Schedule.spaced("3 seconds"),
345 }),
346 Effect.catchTag("Conflict", () =>
347 findOwned(id, env.project, env.region),
348 ),
349 );
350 current = created ?? undefined;
351 }
352
353 if (current === undefined) {
354 return yield* new CachedContentNotResolved({
355 name: output?.name ?? parentOf(env.project, location),
356 });
357 }
358
359 const name = current.name ?? "";
360 const ttlChanged = news.ttl !== undefined && news.ttl.length > 0;
361 const expireChanged =
362 news.expireTime !== undefined &&
363 news.expireTime !== (current.expireTime ?? "");
364
365 if (ttlChanged || expireChanged) {
366 current = yield* aiplatform.patchProjectsLocationsCachedContents({
367 name,
368 updateMask: ttlChanged ? "ttl" : "expire_time",
369 body: {
370 name,
371 ttl: ttlChanged ? news.ttl : undefined,
372 expireTime: expireChanged ? news.expireTime : undefined,
373 },
374 });
375 }
376
377 return toAttrs(current, env.project);
378 }),
379
380 delete: Effect.fn(function* ({ output }) {
381 yield* aiplatform
382 .deleteProjectsLocationsCachedContents({ name: output.name })
383 .pipe(Effect.catchTag("NotFound", () => Effect.void));
384 }),
385 });