Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 37 additions & 2 deletions src/cloudflare/internal/ai-api.ts
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,27 @@ function isFormData(obj: unknown): obj is FormData {
return obj instanceof FormData;
}

/**
* Check if a model ID should be handled by AI Gateway
* AI Gateway models follow the pattern "provider/model-name" (e.g., "xai/grok-voice")
* Workers AI models start with "@cf/" or "@hf/" (e.g., "@cf/meta/llama-3.1-8b-instruct")
*/
function isHandledByAiGateway(modelId: string): boolean {
// Workers AI models
if (modelId.startsWith('@cf/') || modelId.startsWith('@hf/')) {
return false;
}

// Skip processing models (e.g., "skipProcessing:...")
if (modelId.includes(':')) {
return false;
}

// AI Gateway models follow "provider/model-name" pattern
const HANDLED_BY_AIG = /^[a-zA-Z0-9._-]+\/[a-zA-Z0-9._-]+$/;
return HANDLED_BY_AIG.test(modelId);
}

/**
* Find keys in inputs that have a ReadableStream
* */
Expand Down Expand Up @@ -284,8 +305,22 @@ export class Ai {
}

const aiEndpoint = new URL(`${this.#endpointURL}/run`);
aiEndpoint.searchParams.set('version', '3');
aiEndpoint.searchParams.set('body', body);

// AI Gateway models expect inputs as query params, not version/body
if (isHandledByAiGateway(model)) {
aiEndpoint.searchParams.set('model', model);

// Add all inputs as query parameters
for (const [key, value] of Object.entries(inputs)) {
if (value !== null && value !== undefined) {
aiEndpoint.searchParams.set(key, String(value));
}
}
} else {
// Workers AI models use version and body query params
aiEndpoint.searchParams.set('version', '3');
aiEndpoint.searchParams.set('body', body);
}

return await this.#fetcher.fetch(aiEndpoint, fetchOptions);
}
Expand Down
24 changes: 23 additions & 1 deletion src/cloudflare/internal/test/ai/ai-api-test.js
Original file line number Diff line number Diff line change
Expand Up @@ -345,7 +345,7 @@ export const tests = {
}

{
// Test websocket option with basic inputs
// Test websocket option with basic inputs (Workers AI model)
const resp = await env.ai.run(
'@cf/test/websocket',
{ encoding: 'utf8' },
Expand All @@ -366,6 +366,28 @@ export const tests = {
});
}

{
// Test websocket option with AI Gateway model (inputs as query params)
const resp = await env.ai.run(
'xai/grok-voice',
{ text: 'hello', voice: 'alloy' },
{ websocket: true }
);
assert.deepStrictEqual(resp instanceof Response, true);
const respData = await resp.json();
assert.deepStrictEqual(respData, {
inputs: { text: 'hello', voice: 'alloy' },
options: { websocket: true },
requestUrl:
'https://workers-binding.ai/run?model=xai%2Fgrok-voice&text=hello&voice=alloy',
headers: {
'cf-consn-sdk-version': '2.0.0',
'cf-consn-model-id': 'xai/grok-voice',
upgrade: 'websocket',
},
});
}

{
// Test signal option is not included in request body
const controller = new AbortController();
Expand Down
Loading