Refactor Cloudflare AI gateway handling

Updated the Cloudflare AI utility functions and configuration to support a dedicated embedding gateway ID. The changes include a new function to determine the use of the embedding model and adjustments to the workflow files to include the embedding gateway ID as a secret. Tests were also updated to ensure proper error handling when the embedding gateway ID is not set.

Author: kentcdodds

.github/workflows/index-semantic-content.yml

@@ -85,6 +85,8 @@ jobs:
           CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }}
           CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }}
           CLOUDFLARE_AI_GATEWAY_ID: ${{ secrets.CLOUDFLARE_AI_GATEWAY_ID }}
+          CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID:
+            ${{ secrets.CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID }}
           CLOUDFLARE_AI_GATEWAY_AUTH_TOKEN:
             ${{ secrets.CLOUDFLARE_AI_GATEWAY_AUTH_TOKEN }}
           CLOUDFLARE_VECTORIZE_INDEX: ${{ secrets.CLOUDFLARE_VECTORIZE_INDEX }}

.github/workflows/index-semantic-podcasts.yml

@@ -31,6 +31,8 @@ jobs:
           CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }}
           CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }}
           CLOUDFLARE_AI_GATEWAY_ID: ${{ secrets.CLOUDFLARE_AI_GATEWAY_ID }}
+          CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID:
+            ${{ secrets.CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID }}
           CLOUDFLARE_AI_GATEWAY_AUTH_TOKEN:
             ${{ secrets.CLOUDFLARE_AI_GATEWAY_AUTH_TOKEN }}
           CLOUDFLARE_VECTORIZE_INDEX: ${{ secrets.CLOUDFLARE_VECTORIZE_INDEX }}

.github/workflows/index-semantic-youtube.yml

@@ -43,6 +43,8 @@ jobs:
           CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }}
           CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }}
           CLOUDFLARE_AI_GATEWAY_ID: ${{ secrets.CLOUDFLARE_AI_GATEWAY_ID }}
+          CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID:
+            ${{ secrets.CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID }}
           CLOUDFLARE_AI_GATEWAY_AUTH_TOKEN:
             ${{ secrets.CLOUDFLARE_AI_GATEWAY_AUTH_TOKEN }}
           CLOUDFLARE_VECTORIZE_INDEX: ${{ secrets.CLOUDFLARE_VECTORIZE_INDEX }}

app/utils/__tests__/cloudflare-ai-utils.server.test.ts

@@ -0,0 +1,42 @@
+import { expect, test } from 'vitest'
+import { setEnv } from '#tests/env-disposable.ts'
+import { getWorkersAiRunUrl } from '../cloudflare-ai-utils.server.ts'
+
+test('getWorkersAiRunUrl routes embeddinggemma requests through CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID', () => {
+	using ignoredEnv = setEnv({
+		CLOUDFLARE_ACCOUNT_ID: 'cf-account',
+		CLOUDFLARE_AI_GATEWAY_ID: 'runtime-gateway',
+		CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID: 'embedding-gateway',
+	})
+
+	const url = getWorkersAiRunUrl('@cf/google/embeddinggemma-300m')
+	expect(url).toContain('/embedding-gateway/')
+	expect(url).not.toContain('/runtime-gateway/')
+})
+
+test('getWorkersAiRunUrl keeps non-embedding models on CLOUDFLARE_AI_GATEWAY_ID', () => {
+	using ignoredEnv = setEnv({
+		CLOUDFLARE_ACCOUNT_ID: 'cf-account',
+		CLOUDFLARE_AI_GATEWAY_ID: 'runtime-gateway',
+		CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID: 'embedding-gateway',
+	})
+
+	const url = getWorkersAiRunUrl('@cf/openai/whisper-large-v3-turbo')
+	expect(url).toContain('/runtime-gateway/')
+	expect(url).not.toContain('/embedding-gateway/')
+})
+
+test('getWorkersAiRunUrl prefers explicit gatewayId overrides', () => {
+	using ignoredEnv = setEnv({
+		CLOUDFLARE_ACCOUNT_ID: 'cf-account',
+		CLOUDFLARE_AI_GATEWAY_ID: 'runtime-gateway',
+		CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID: 'embedding-gateway',
+	})
+
+	const url = getWorkersAiRunUrl({
+		model: '@cf/google/embeddinggemma-300m',
+		gatewayId: 'explicit-gateway',
+	})
+	expect(url).toContain('/explicit-gateway/')
+	expect(url).not.toContain('/embedding-gateway/')
+})

app/utils/cloudflare-ai-utils.server.ts

@@ -6,14 +6,24 @@ type WorkersAiRunUrlOptions = {
 	gatewayId?: string
 }
 
+const EMBEDDING_GEMMA_MODEL_SEGMENT = 'embeddinggemma-300m'
+
+function usesEmbeddingGemmaModel(model: string) {
+	return model.includes(EMBEDDING_GEMMA_MODEL_SEGMENT)
+}
+
 export function getWorkersAiRunUrl(options: string | WorkersAiRunUrlOptions) {
 	const { model, accountId, gatewayId } =
 		typeof options === 'string' ? { model: options } : options
 	// Cloudflare's REST route expects the model as path segments (with `/`), so do
 	// not URL-encode the model string (encoding can yield "No route for that URI").
 	const env = getEnv()
 	const resolvedAccountId = accountId ?? env.CLOUDFLARE_ACCOUNT_ID
-	const resolvedGatewayId = gatewayId ?? env.CLOUDFLARE_AI_GATEWAY_ID
+	const resolvedGatewayId =
+		gatewayId ??
+		(usesEmbeddingGemmaModel(model)
+			? env.CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID
+			: env.CLOUDFLARE_AI_GATEWAY_ID)
 	return `https://gateway.ai.cloudflare.com/v1/${resolvedAccountId}/${resolvedGatewayId}/workers-ai/${model}`
 }
 

other/semantic-search/__tests__/cloudflare-config.test.ts

@@ -16,7 +16,7 @@ test('getCloudflareConfig prefers embedding gateway for indexing when configured
 	expect(config.gatewayId).toBe('indexing-gateway')
 })
 
-test('getCloudflareConfig falls back to regular gateway when embedding override is unset', () => {
+test('getCloudflareConfig throws when embedding gateway is unset', () => {
 	using ignoredEnv = setEnv({
 		CLOUDFLARE_ACCOUNT_ID: 'cf-account',
 		CLOUDFLARE_API_TOKEN: 'cf-token',
@@ -26,6 +26,7 @@ test('getCloudflareConfig falls back to regular gateway when embedding override
 		CLOUDFLARE_VECTORIZE_INDEX: 'vector-index',
 	})
 
-	const config = getCloudflareConfig()
-	expect(config.gatewayId).toBe('runtime-gateway')
+	expect(() => getCloudflareConfig()).toThrow(
+		'Missing required env var: CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID',
+	)
 })

other/semantic-search/cloudflare.ts

@@ -14,15 +14,10 @@ function getRequiredEnv(name: string) {
 }
 
 export function getCloudflareConfig() {
-	const defaultGatewayId = getRequiredEnv('CLOUDFLARE_AI_GATEWAY_ID')
-	const embeddingGatewayId =
-		process.env.CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID?.trim() || defaultGatewayId
-
 	return {
 		accountId: getRequiredEnv('CLOUDFLARE_ACCOUNT_ID'),
 		apiToken: getRequiredEnv('CLOUDFLARE_API_TOKEN'),
-		// Embedding jobs can use a dedicated gateway without guardrails.
-		gatewayId: embeddingGatewayId,
+		gatewayId: getRequiredEnv('CLOUDFLARE_AI_EMBEDDING_GATEWAY_ID'),
 		gatewayAuthToken: getRequiredEnv('CLOUDFLARE_AI_GATEWAY_AUTH_TOKEN'),
 		vectorizeIndex: getRequiredEnv('CLOUDFLARE_VECTORIZE_INDEX'),
 		embeddingModel: