1import * as aiGateway from "@distilled.cloud/cloudflare/ai-gateway";
2import * as Effect from "effect/Effect";
3import * as Predicate from "effect/Predicate";
4import * as Stream from "effect/Stream";
5import { deepEqual, isResolved } from "../../Diff.ts";
6import { createPhysicalName } from "../../PhysicalName.ts";
7import * as Provider from "../../Provider.ts";
8import { Resource } from "../../Resource.ts";
9import { CloudflareEnvironment } from "../CloudflareEnvironment.ts";
10import type { Providers } from "../Providers.ts";
11
12const TypeId = "Cloudflare.AI.Dataset" as const;
13type TypeId = typeof TypeId;
14
15/**
16 * Log property a dataset filter can match on.
17 */
18export type DatasetFilterKey =
19 | "created_at"
20 | "request_content_type"
21 | "response_content_type"
22 | "success"
23 | "cached"
24 | "provider"
25 | "model"
26 | "cost"
27 | "tokens"
28 | "tokens_in"
29 | "tokens_out"
30 | "duration"
31 | "feedback";
32
33/**
34 * Comparison operator applied to a dataset filter.
35 */
36export type DatasetFilterOperator = "eq" | "contains" | "lt" | "gt";
37
38/**
39 * A single saved log filter on an AI Gateway dataset.
40 */
41export type DatasetFilter = {
42 /**
43 * Log property to match on (e.g. `success`, `model`, `provider`).
44 */
45 key: DatasetFilterKey;
46 /**
47 * Comparison operator.
48 */
49 operator: DatasetFilterOperator;
50 /**
51 * Values to compare against. Multiple values act as an OR.
52 */
53 value: (string | number | boolean)[];
54};
55
56export type DatasetProps = {
57 /**
58 * The AI Gateway the dataset belongs to. Changing the gateway triggers a
59 * replacement.
60 */
61 gatewayId: string;
62 /**
63 * Human readable dataset name. If omitted, a unique name is generated from
64 * the app, stage, and logical ID.
65 * @default ${app}-${stage}-${id}
66 */
67 name?: string;
68 /**
69 * Whether the dataset is enabled (actively collecting matching logs).
70 * @default true
71 */
72 enable?: boolean;
73 /**
74 * Saved log filters defining which gateway logs the dataset captures.
75 * An empty array captures all logs.
76 */
77 filters: DatasetFilter[];
78};
79
80export type DatasetAttributes = {
81 /**
82 * Server-generated dataset identifier. Stable across updates.
83 */
84 datasetId: string;
85 /**
86 * The Cloudflare account the dataset belongs to.
87 */
88 accountId: string;
89 /**
90 * The AI Gateway the dataset belongs to.
91 */
92 gatewayId: string;
93 /**
94 * Human readable dataset name.
95 */
96 name: string;
97 /**
98 * Whether the dataset is enabled.
99 */
100 enable: boolean;
101 /**
102 * Saved log filters.
103 */
104 filters: DatasetFilter[];
105 /**
106 * When the dataset was created.
107 */
108 createdAt: string;
109 /**
110 * When the dataset was last modified.
111 */
112 modifiedAt: string;
113};
114
115export type Dataset = Resource<
116 TypeId,
117 DatasetProps,
118 DatasetAttributes,
119 never,
120 Providers
121>;
122
123/**
124 * A saved log filter ("dataset") on a Cloudflare.AI. Gateway.
125 *
126 * Datasets capture a slice of the gateway's request logs (filtered by
127 * provider, model, success, cost, tokens, etc.) and serve as the input to AI
128 * Gateway evaluations. Name, enablement, and filters are all mutable in
129 * place; only moving the dataset to a different gateway forces a replacement.
130 * ### Creating a Dataset
131 * **Example:** Capture successful requests
132 * ```typescript
133 * const gateway = yield* Cloudflare.AI.Gateway("Gateway");
134 *
135 * const dataset = yield* Cloudflare.AI.Dataset("SuccessLogs", {
136 * gatewayId: gateway.gatewayId,
137 * filters: [{ key: "success", operator: "eq", value: [true] }],
138 * });
139 * ```
140 *
141 * **Example:** Capture logs for a specific model
142 * ```typescript
143 * const dataset = yield* Cloudflare.AI.Dataset("LlamaLogs", {
144 * gatewayId: gateway.gatewayId,
145 * name: "llama-traffic",
146 * filters: [
147 * { key: "provider", operator: "eq", value: ["workers-ai"] },
148 * { key: "model", operator: "contains", value: ["llama"] },
149 * ],
150 * });
151 * ```
152 *
153 * ### Updating a Dataset
154 * **Example:** Disable collection without deleting
155 * ```typescript
156 * const dataset = yield* Cloudflare.AI.Dataset("SuccessLogs", {
157 * gatewayId: gateway.gatewayId,
158 * enable: false,
159 * filters: [{ key: "success", operator: "eq", value: [true] }],
160 * });
161 * ```
162 *
163 * @see https://developers.cloudflare.com/ai-gateway/evaluations/set-up-evaluations/
164 *
165 * @resource
166 * @product AI Gateway
167 * @category AI
168 */
169export const Dataset = Resource<Dataset>(TypeId, {
170 aliases: ["Cloudflare.AiGateway.Dataset"],
171});
172
173/**
174 * Returns true if the given value is a Dataset resource.
175 */
176export const isDataset = (value: unknown): value is Dataset =>
177 Predicate.hasProperty(value, "Type") && value.Type === TypeId;
178
179export const DatasetProvider = () =>
180 Provider.succeed(Dataset, {
181 stables: ["datasetId", "accountId", "gatewayId", "createdAt"],
182 diff: Effect.fn(function* ({ olds, news, output }) {
183 if (!isResolved(news)) return undefined;
184 const { accountId } = yield* yield* CloudflareEnvironment;
185 if ((output?.accountId ?? accountId) !== accountId) {
186 return { action: "replace" } as const;
187 }
188 // The gateway is a path parameter — a dataset cannot move between
189 // gateways in place. By diff time both sides are resolved strings.
190 const oldGatewayId = output?.gatewayId ?? olds?.gatewayId;
191 if (
192 typeof oldGatewayId === "string" &&
193 typeof news.gatewayId === "string" &&
194 oldGatewayId !== news.gatewayId
195 ) {
196 return { action: "replace" } as const;
197 }
198 return undefined;
199 }),
200 read: Effect.fn(function* ({ id, olds, output }) {
201 const { accountId } = yield* yield* CloudflareEnvironment;
202 const acct = output?.accountId ?? accountId;
203 const gatewayId =
204 output?.gatewayId ?? (olds?.gatewayId as string | undefined);
205 if (gatewayId === undefined) return undefined;
206
207 if (output?.datasetId) {
208 const observed = yield* getDataset(acct, gatewayId, output.datasetId);
209 return observed ? toAttributes(observed, acct) : undefined;
210 }
211 // Cold read — recover from lost state by matching the deterministic
212 // physical name. Names are not unique on Cloudflare's side; an exact
213 // match on our generated/explicit name is the best identity we have.
214 const name = yield* createDatasetName(id, olds?.name);
215 const match = yield* findByName(acct, gatewayId, name);
216 return match ? toAttributes(match, acct) : undefined;
217 }),
218 reconcile: Effect.fn(function* ({ id, news, output }) {
219 const { accountId } = yield* yield* CloudflareEnvironment;
220 const gatewayId = news.gatewayId as string;
221 const name = yield* createDatasetName(id, news.name);
222 const desired = {
223 name,
224 enable: news.enable ?? true,
225 filters: news.filters ?? [],
226 };
227
228 // Observe — the datasetId cached on `output` is a hint, not a
229 // guarantee: a missing dataset falls through and we recreate.
230 let observed = output?.datasetId
231 ? yield* getDataset(
232 output.accountId ?? accountId,
233 gatewayId,
234 output.datasetId,
235 )
236 : undefined;
237
238 if (!observed) {
239 // Ensure — greenfield (or out-of-band delete). Cloudflare enforces
240 // name uniqueness per gateway, so a name conflict (e.g. a leftover
241 // from an interrupted run or a create race) is converged by
242 // adopting the existing dataset and syncing it to the desired
243 // shape below instead of failing.
244 const created = yield* aiGateway
245 .createDataset({
246 accountId,
247 gatewayId,
248 ...desired,
249 })
250 .pipe(
251 Effect.catchTag("DatasetNameAlreadyExists", (originalError) =>
252 findByName(accountId, gatewayId, name).pipe(
253 Effect.flatMap((match) =>
254 match ? Effect.succeed(match) : Effect.fail(originalError),
255 ),
256 ),
257 ),
258 );
259 observed = created;
260 }
261
262 // Sync — diff observed cloud state against desired; the update API is
263 // a PUT that takes the full body, so send everything, but skip the
264 // call entirely on a no-op.
265 const observedShape = {
266 name: observed.name,
267 enable: observed.enable,
268 filters: normalizeFilters(observed.filters),
269 };
270 if (deepEqual(observedShape, desired)) {
271 return toAttributes(observed, accountId);
272 }
273 const updated = yield* aiGateway.updateDataset({
274 accountId,
275 gatewayId,
276 id: observed.id,
277 ...desired,
278 });
279 return toAttributes(updated, accountId);
280 }),
281 delete: Effect.fn(function* ({ output }) {
282 yield* aiGateway
283 .deleteDataset({
284 accountId: output.accountId,
285 gatewayId: output.gatewayId,
286 id: output.datasetId,
287 })
288 // Cloudflare reports both a missing dataset and a missing gateway
289 // with code 7002 on this endpoint — either way it's already gone.
290 .pipe(Effect.catchTag("DatasetNotFound", () => Effect.void));
291 }),
292 list: Effect.fn(function* () {
293 const { accountId } = yield* yield* CloudflareEnvironment;
294 // Datasets are scoped under a gateway and there is no account-wide
295 // dataset list, so fan out: enumerate every account gateway, then
296 // exhaustively list each gateway's datasets.
297 const gateways = yield* aiGateway.listAiGateways
298 .pages({ accountId })
299 .pipe(
300 Stream.runCollect,
301 Effect.map((chunk) =>
302 Array.from(chunk).flatMap((page) => page.result ?? []),
303 ),
304 );
305 const rows = yield* Effect.forEach(
306 gateways,
307 (gateway) =>
308 aiGateway.listDatasets
309 .pages({ accountId, gatewayId: gateway.id })
310 .pipe(
311 Stream.runCollect,
312 Effect.map((chunk) =>
313 Array.from(chunk).flatMap((page) =>
314 (page.result ?? []).map((d) => toAttributes(d, accountId)),
315 ),
316 ),
317 // A gateway removed between enumeration and its dataset list
318 // is gone — skip it rather than failing the whole listing.
319 Effect.catchTag("GatewayNotFound", () =>
320 Effect.succeed([] as DatasetAttributes[]),
321 ),
322 ),
323 { concurrency: 10 },
324 );
325 return rows.flat();
326 }),
327 });
328
329/**
330 * Read a dataset by id, mapping "gone" (`DatasetNotFound`, Cloudflare error
331 * code 7002 — which also covers a deleted parent gateway) to `undefined`.
332 */
333const getDataset = (accountId: string, gatewayId: string, id: string) =>
334 aiGateway
335 .getDataset({ accountId, gatewayId, id })
336 .pipe(Effect.catchTag("DatasetNotFound", () => Effect.succeed(undefined)));
337
338/**
339 * Find a dataset by exact name. Cloudflare's `name` query is a fuzzy match,
340 * so re-check exactly client-side. If several datasets carry the same name,
341 * pick the oldest for determinism. A missing parent gateway means the
342 * dataset is gone too.
343 */
344const findByName = (accountId: string, gatewayId: string, name: string) =>
345 aiGateway.listDatasets
346 .items({ accountId, gatewayId, name, perPage: 50 })
347 .pipe(
348 Stream.filter((d) => d.name === name),
349 Stream.runCollect,
350 Effect.map((chunk) =>
351 Array.from(chunk)
352 .sort((a, b) => a.createdAt.localeCompare(b.createdAt))
353 .at(0),
354 ),
355 Effect.catchTag("GatewayNotFound", () => Effect.succeed(undefined)),
356 );
357
358const createDatasetName = (id: string, name: string | undefined) =>
359 Effect.gen(function* () {
360 return name ?? (yield* createPhysicalName({ id, lowercase: true }));
361 });
362
363/**
364 * Strip readonly markers / open-union widening from the wire filters so they
365 * diff structurally against the user's props.
366 */
367const normalizeFilters = (
368 filters: readonly {
369 key: string;
370 operator: string;
371 value: readonly (string | number | boolean)[];
372 }[],
373): DatasetFilter[] =>
374 filters.map((f) => ({
375 key: f.key as DatasetFilterKey,
376 operator: f.operator as DatasetFilterOperator,
377 value: [...f.value],
378 }));
379
380const toAttributes = (
381 dataset:
382 | aiGateway.GetDatasetResponse
383 | aiGateway.CreateDatasetResponse
384 | aiGateway.UpdateDatasetResponse,
385 accountId: string,
386): DatasetAttributes => ({
387 datasetId: dataset.id,
388 accountId,
389 gatewayId: dataset.gatewayId,
390 name: dataset.name,
391 enable: dataset.enable,
392 filters: normalizeFilters(dataset.filters),
393 createdAt: dataset.createdAt,
394 modifiedAt: dataset.modifiedAt,
395});