| 1 | import * as vectorize from "@distilled.cloud/cloudflare/vectorize"; |
| 2 | import * as Effect from "effect/Effect"; |
| 3 | import * as Stream from "effect/Stream"; |
| 4 | |
| 5 | import { isResolved } from "../../Diff.ts"; |
| 6 | import { createPhysicalName } from "../../PhysicalName.ts"; |
| 7 | import * as Provider from "../../Provider.ts"; |
| 8 | import { isResourceOfType, Resource } from "../../Resource.ts"; |
| 9 | import { CloudflareEnvironment } from "../CloudflareEnvironment.ts"; |
| 10 | import type { Providers } from "../Providers.ts"; |
| 11 | |
| 12 | const TypeId = "Cloudflare.VectorizeIndex" as const; |
| 13 | type TypeId = typeof TypeId; |
| 14 | |
| 15 | export type DistanceMetric = "cosine" | "euclidean" | "dot-product"; |
| 16 | |
| 17 | export type Preset = |
| 18 | | "@cf/baai/bge-small-en-v1.5" |
| 19 | | "@cf/baai/bge-base-en-v1.5" |
| 20 | | "@cf/baai/bge-large-en-v1.5" |
| 21 | | "openai/text-embedding-ada-002" |
| 22 | | "cohere/embed-multilingual-v2.0" |
| 23 | // Keep the union open so new Cloudflare presets aren't blocked by stale types. |
| 24 | | (string & {}); |
| 25 | |
| 26 | const DEFAULT_METRIC: DistanceMetric = "cosine"; |
| 27 | |
| 28 | export type IndexProps = { |
| 29 | /** |
| 30 | * Name of the index. If omitted, a unique name will be generated. |
| 31 | * Must be lowercase alphanumeric with hyphens. Changing it triggers a |
| 32 | * replacement. |
| 33 | * @default ${app}-${stage}-${id} |
| 34 | */ |
| 35 | name?: string; |
| 36 | /** |
| 37 | * Number of dimensions for each vector stored in the index. Required |
| 38 | * unless `preset` is provided. Cannot be changed after creation — |
| 39 | * updating this property triggers a replacement. |
| 40 | */ |
| 41 | dimensions?: number; |
| 42 | /** |
| 43 | * Distance metric used for similarity search. Cannot be changed after |
| 44 | * creation — updating this property triggers a replacement. |
| 45 | * @default "cosine" |
| 46 | */ |
| 47 | metric?: DistanceMetric; |
| 48 | /** |
| 49 | * A managed embedding model preset that fixes `dimensions` and `metric` |
| 50 | * to match the named model. Mutually exclusive with `dimensions`/`metric`. |
| 51 | * Cannot be changed after creation — updating this triggers a replacement. |
| 52 | */ |
| 53 | preset?: Preset; |
| 54 | /** |
| 55 | * Human-readable description of the index. Vectorize has no update API, |
| 56 | * so changing the description triggers a replacement. |
| 57 | */ |
| 58 | description?: string; |
| 59 | }; |
| 60 | |
| 61 | export type IndexAttributes = { |
| 62 | indexName: string; |
| 63 | dimensions: number | undefined; |
| 64 | metric: DistanceMetric | undefined; |
| 65 | description: string | undefined; |
| 66 | accountId: string; |
| 67 | createdOn: string | undefined; |
| 68 | modifiedOn: string | undefined; |
| 69 | }; |
| 70 | |
| 71 | export type Index = Resource< |
| 72 | TypeId, |
| 73 | IndexProps, |
| 74 | IndexAttributes, |
| 75 | never, |
| 76 | Providers |
| 77 | >; |
| 78 | |
| 79 | /** |
| 80 | * A Cloudflare Vectorize index for storing and querying vector embeddings. |
| 81 | * |
| 82 | * Vectorize is a globally distributed vector database. Create an index as a |
| 83 | * resource, then bind it to a Worker to insert, upsert, and query vectors. |
| 84 | * |
| 85 | * A Vectorize index is identified by its name and is immutable: its |
| 86 | * dimensions, metric, preset, and description are all fixed at creation. |
| 87 | * Changing any of them triggers a replacement. |
| 88 | * ### Creating an Index |
| 89 | * **Example:** Index with explicit dimensions and metric |
| 90 | * ```typescript |
| 91 | * const index = yield* Cloudflare.Vectorize.Index("my-index", { |
| 92 | * dimensions: 768, |
| 93 | * metric: "cosine", |
| 94 | * }); |
| 95 | * ``` |
| 96 | * |
| 97 | * **Example:** Index from a managed embedding model preset |
| 98 | * A preset fixes the dimensions and metric to match the named model. |
| 99 | * ```typescript |
| 100 | * const index = yield* Cloudflare.Vectorize.Index("my-index", { |
| 101 | * preset: "@cf/baai/bge-base-en-v1.5", |
| 102 | * }); |
| 103 | * ``` |
| 104 | * |
| 105 | * **Example:** Index with a description |
| 106 | * ```typescript |
| 107 | * const index = yield* Cloudflare.Vectorize.Index("my-index", { |
| 108 | * dimensions: 1536, |
| 109 | * metric: "euclidean", |
| 110 | * description: "Product catalog embeddings", |
| 111 | * }); |
| 112 | * ``` |
| 113 | * |
| 114 | * ### Binding to a Worker |
| 115 | * **Example:** Querying an index inside a Worker |
| 116 | * ```typescript |
| 117 | * const index = yield* Cloudflare.Vectorize.SearchIndex(MyIndex); |
| 118 | * |
| 119 | * // Insert vectors |
| 120 | * yield* index.upsert([ |
| 121 | * { id: "1", values: [0.1, 0.2, 0.3], metadata: { title: "doc" } }, |
| 122 | * ]); |
| 123 | * |
| 124 | * // Query the nearest neighbors |
| 125 | * const matches = yield* index.query([0.1, 0.2, 0.3], { topK: 5 }); |
| 126 | * ``` |
| 127 | * |
| 128 | * @see https://developers.cloudflare.com/vectorize/ |
| 129 | * |
| 130 | * @resource |
| 131 | * @product Vectorize |
| 132 | * @category AI |
| 133 | */ |
| 134 | export const Index = Resource<Index>(TypeId); |
| 135 | |
| 136 | /** |
| 137 | * Returns true if the given value is a Vectorize Index resource. |
| 138 | */ |
| 139 | export const isIndex = (value: unknown): value is Index => |
| 140 | isResourceOfType(value, TypeId); |
| 141 | |
| 142 | export const IndexProvider = () => |
| 143 | Provider.succeed(Index, { |
| 144 | stables: ["indexName", "accountId"], |
| 145 | diff: Effect.fn(function* ({ id, olds = {}, news = {}, output }) { |
| 146 | const { accountId } = yield* yield* CloudflareEnvironment; |
| 147 | if (!isResolved(news)) return undefined; |
| 148 | if ((output?.accountId ?? accountId) !== accountId) { |
| 149 | return { action: "replace" } as const; |
| 150 | } |
| 151 | const oldName = |
| 152 | output?.indexName ?? (yield* createIndexName(id, olds.name)); |
| 153 | // Auto-generated names are engine-owned: the deployed name stays |
| 154 | // authoritative even if the generator would name this id differently |
| 155 | // today. Only an explicit user-provided name can force a replace. |
| 156 | const name = news.name ?? oldName; |
| 157 | if ( |
| 158 | oldName !== name || |
| 159 | (news.preset ?? undefined) !== (olds.preset ?? undefined) || |
| 160 | (news.dimensions ?? undefined) !== (olds.dimensions ?? undefined) || |
| 161 | (news.metric ?? DEFAULT_METRIC) !== (olds.metric ?? DEFAULT_METRIC) || |
| 162 | (news.description ?? undefined) !== (olds.description ?? undefined) |
| 163 | ) { |
| 164 | return { action: "replace" } as const; |
| 165 | } |
| 166 | return undefined; |
| 167 | }), |
| 168 | read: Effect.fn(function* ({ id, output, olds }) { |
| 169 | const { accountId } = yield* yield* CloudflareEnvironment; |
| 170 | const acct = output?.accountId ?? accountId; |
| 171 | const name = |
| 172 | output?.indexName ?? (yield* createIndexName(id, olds?.name)); |
| 173 | return yield* vectorize |
| 174 | .getIndex({ accountId: acct, indexName: name }) |
| 175 | .pipe( |
| 176 | Effect.map((index) => toAttributes(index, name, acct)), |
| 177 | Effect.catchTag(["NotFound", "Gone"], () => |
| 178 | Effect.succeed(undefined), |
| 179 | ), |
| 180 | ); |
| 181 | }), |
| 182 | reconcile: Effect.fn(function* ({ id, news = {}, output }) { |
| 183 | const { accountId } = yield* yield* CloudflareEnvironment; |
| 184 | // Prefer the deployed name: regenerating would target a different |
| 185 | // index if the generator's output for this id ever drifts. |
| 186 | const indexName = |
| 187 | output?.indexName ?? (yield* createIndexName(id, news.name)); |
| 188 | |
| 189 | // Observe — read the live index by name. The name is the stable |
| 190 | // identifier; fall back through a NotFound to the create path so |
| 191 | // we recover from out-of-band deletes or partial state-persistence. |
| 192 | let observed = yield* vectorize |
| 193 | .getIndex({ |
| 194 | accountId, |
| 195 | indexName, |
| 196 | }) |
| 197 | .pipe( |
| 198 | Effect.catchTag(["NotFound", "Gone"], () => |
| 199 | Effect.succeed(undefined), |
| 200 | ), |
| 201 | ); |
| 202 | |
| 203 | // Ensure — create if missing. Cloudflare returns 409 Conflict when |
| 204 | // an index with the same name already exists; tolerate the race by |
| 205 | // re-reading it. |
| 206 | if (!observed) { |
| 207 | observed = yield* vectorize |
| 208 | .createIndex({ |
| 209 | accountId, |
| 210 | name: indexName, |
| 211 | config: buildConfig(news), |
| 212 | description: news.description, |
| 213 | }) |
| 214 | .pipe( |
| 215 | Effect.catchTag("IndexAlreadyExists", () => |
| 216 | vectorize.getIndex({ accountId, indexName }), |
| 217 | ), |
| 218 | ); |
| 219 | } |
| 220 | |
| 221 | return toAttributes(observed, indexName, accountId); |
| 222 | }), |
| 223 | delete: Effect.fn(function* ({ output }) { |
| 224 | yield* vectorize |
| 225 | .deleteIndex({ |
| 226 | accountId: output.accountId, |
| 227 | indexName: output.indexName, |
| 228 | }) |
| 229 | .pipe(Effect.catchTag(["NotFound", "Gone"], () => Effect.void)); |
| 230 | }), |
| 231 | list: Effect.fn(function* () { |
| 232 | const { accountId } = yield* yield* CloudflareEnvironment; |
| 233 | return yield* vectorize.listIndexes.pages({ accountId }).pipe( |
| 234 | Stream.runCollect, |
| 235 | Effect.map((chunk) => |
| 236 | Array.from(chunk).flatMap((page) => |
| 237 | (page.result ?? []) |
| 238 | .filter( |
| 239 | ( |
| 240 | index, |
| 241 | ): index is (typeof page.result)[number] & { |
| 242 | name: string; |
| 243 | } => index.name != null, |
| 244 | ) |
| 245 | .map((index) => toAttributes(index, index.name, accountId)), |
| 246 | ), |
| 247 | ), |
| 248 | ); |
| 249 | }), |
| 250 | }); |
| 251 | |
| 252 | const toAttributes = ( |
| 253 | index: vectorize.GetIndexResponse | vectorize.CreateIndexResponse, |
| 254 | name: string, |
| 255 | accountId: string, |
| 256 | ): IndexAttributes => ({ |
| 257 | indexName: index.name ?? name, |
| 258 | dimensions: index.config?.dimensions, |
| 259 | // Distilled widened generated string enums to open unions (`string & {}`). |
| 260 | metric: index.config?.metric as DistanceMetric | undefined, |
| 261 | description: index.description ?? undefined, |
| 262 | accountId, |
| 263 | createdOn: index.createdOn ?? undefined, |
| 264 | modifiedOn: index.modifiedOn ?? undefined, |
| 265 | }); |
| 266 | |
| 267 | const createIndexName = (id: string, name: string | undefined) => |
| 268 | Effect.gen(function* () { |
| 269 | return name ?? (yield* createPhysicalName({ id, lowercase: true })); |
| 270 | }); |
| 271 | |
| 272 | const buildConfig = ( |
| 273 | news: IndexProps, |
| 274 | ): vectorize.CreateIndexRequest["config"] => |
| 275 | news.preset !== undefined |
| 276 | ? // `Preset` is intentionally open (`| (string & {})`) so |
| 277 | // new Cloudflare presets aren't blocked by stale types. The |
| 278 | // distilled type is the strict-at-release-time union; cast |
| 279 | // through for the API call. |
| 280 | ({ |
| 281 | preset: news.preset as never, |
| 282 | } as vectorize.CreateIndexRequest["config"]) |
| 283 | : { |
| 284 | dimensions: news.dimensions!, |
| 285 | metric: news.metric ?? DEFAULT_METRIC, |
| 286 | }; |