Add returnRawResponse to AI binding options

cloudflare · Jan 29, 2025 · 101e0c0 · 101e0c0
1 parent 270a365
commit 101e0c0
Show file tree

Hide file tree

Showing 23 changed files with 225 additions and 64 deletions.
diff --git a/src/cloudflare/internal/ai-api.ts b/src/cloudflare/internal/ai-api.ts
@@ -23,7 +23,8 @@ export type SessionOptions = {
 
 export type AiOptions = {
   gateway?: GatewayOptions;
-
+  /** If true it will return a Response object */
+  returnRawResponse?: boolean;
   prefix?: string;
   extraHeaders?: object;
   /*
@@ -142,6 +143,10 @@ export class Ai {
     this.aiGatewayLogId = res.headers.get('cf-aig-log-id');
     this.lastRequestHttpStatusCode = res.status;
 
+    if (this.options.returnRawResponse) {
+      return res;
+    }
+
     if (!res.ok) {
       throw await this._parseError(res);
     }

diff --git a/src/cloudflare/internal/test/ai/ai-api-test.js b/src/cloudflare/internal/test/ai/ai-api-test.js
@@ -196,5 +196,16 @@ export const tests = {
         },
       ]);
     }
+
+    {
+      // Test `returnRawResponse` option is returning a Response object
+      const resp = await env.ai.run(
+        'rawInputs',
+        { prompt: 'test' },
+        { returnRawResponse: true }
+      );
+
+      assert.ok(resp instanceof Response);
+    }
   },
 };
diff --git a/types/defines/ai.d.ts b/types/defines/ai.d.ts
@@ -686,6 +686,7 @@ export interface AiModels {
 }
 export type AiOptions = {
   gateway?: GatewayOptions;
+  returnRawResponse?: boolean;
   prefix?: string;
   extraHeaders?: object;
 };
@@ -720,10 +721,14 @@ export type AiModelListType = Record<string, any>;
 export declare abstract class Ai<AiModelList extends AiModelListType = AiModels> {
   aiGatewayLogId: string | null;
   gateway(gatewayId: string): AiGateway;
-  run<Name extends keyof AiModelList>(
+  run<Name extends keyof AiModelList, Options extends AiOptions>(
     model: Name,
     inputs: AiModelList[Name]["inputs"],
-    options?: AiOptions,
-  ): Promise<AiModelList[Name]["postProcessedOutputs"]>;
+    options?: Options
+  ): Promise<
+    Options extends { returnRawResponse: true }
+      ? Response
+      : AiModelList[Name]["postProcessedOutputs"]
+  >;
   public models(params?: AiModelsSearchParams): Promise<AiModelsSearchObject[]>;
 }
diff --git a/types/generated-snapshot/2021-11-03/index.d.ts b/types/generated-snapshot/2021-11-03/index.d.ts
@@ -4110,6 +4110,7 @@ interface AiModels {
 }
 type AiOptions = {
   gateway?: GatewayOptions;
+  returnRawResponse?: boolean;
   prefix?: string;
   extraHeaders?: object;
 };
@@ -4144,11 +4145,17 @@ type AiModelListType = Record<string, any>;
 declare abstract class Ai<AiModelList extends AiModelListType = AiModels> {
   aiGatewayLogId: string | null;
   gateway(gatewayId: string): AiGateway;
-  run<Name extends keyof AiModelList>(
+  run<Name extends keyof AiModelList, Options extends AiOptions>(
     model: Name,
     inputs: AiModelList[Name]["inputs"],
-    options?: AiOptions,
-  ): Promise<AiModelList[Name]["postProcessedOutputs"]>;
+    options?: Options,
+  ): Promise<
+    Options extends {
+      returnRawResponse: true;
+    }
+      ? Response
+      : AiModelList[Name]["postProcessedOutputs"]
+  >;
   public models(params?: AiModelsSearchParams): Promise<AiModelsSearchObject[]>;
 }
 type GatewayOptions = {

diff --git a/types/generated-snapshot/2021-11-03/index.ts b/types/generated-snapshot/2021-11-03/index.ts
@@ -4122,6 +4122,7 @@ export interface AiModels {
 }
 export type AiOptions = {
   gateway?: GatewayOptions;
+  returnRawResponse?: boolean;
   prefix?: string;
   extraHeaders?: object;
 };
@@ -4158,11 +4159,17 @@ export declare abstract class Ai<
 > {
   aiGatewayLogId: string | null;
   gateway(gatewayId: string): AiGateway;
-  run<Name extends keyof AiModelList>(
+  run<Name extends keyof AiModelList, Options extends AiOptions>(
     model: Name,
     inputs: AiModelList[Name]["inputs"],
-    options?: AiOptions,
-  ): Promise<AiModelList[Name]["postProcessedOutputs"]>;
+    options?: Options,
+  ): Promise<
+    Options extends {
+      returnRawResponse: true;
+    }
+      ? Response
+      : AiModelList[Name]["postProcessedOutputs"]
+  >;
   public models(params?: AiModelsSearchParams): Promise<AiModelsSearchObject[]>;
 }
 export type GatewayOptions = {

diff --git a/types/generated-snapshot/2022-01-31/index.d.ts b/types/generated-snapshot/2022-01-31/index.d.ts
@@ -4136,6 +4136,7 @@ interface AiModels {
 }
 type AiOptions = {
   gateway?: GatewayOptions;
+  returnRawResponse?: boolean;
   prefix?: string;
   extraHeaders?: object;
 };
@@ -4170,11 +4171,17 @@ type AiModelListType = Record<string, any>;
 declare abstract class Ai<AiModelList extends AiModelListType = AiModels> {
   aiGatewayLogId: string | null;
   gateway(gatewayId: string): AiGateway;
-  run<Name extends keyof AiModelList>(
+  run<Name extends keyof AiModelList, Options extends AiOptions>(
     model: Name,
     inputs: AiModelList[Name]["inputs"],
-    options?: AiOptions,
-  ): Promise<AiModelList[Name]["postProcessedOutputs"]>;
+    options?: Options,
+  ): Promise<
+    Options extends {
+      returnRawResponse: true;
+    }
+      ? Response
+      : AiModelList[Name]["postProcessedOutputs"]
+  >;
   public models(params?: AiModelsSearchParams): Promise<AiModelsSearchObject[]>;
 }
 type GatewayOptions = {

diff --git a/types/generated-snapshot/2022-01-31/index.ts b/types/generated-snapshot/2022-01-31/index.ts
@@ -4148,6 +4148,7 @@ export interface AiModels {
 }
 export type AiOptions = {
   gateway?: GatewayOptions;
+  returnRawResponse?: boolean;
   prefix?: string;
   extraHeaders?: object;
 };
@@ -4184,11 +4185,17 @@ export declare abstract class Ai<
 > {
   aiGatewayLogId: string | null;
   gateway(gatewayId: string): AiGateway;
-  run<Name extends keyof AiModelList>(
+  run<Name extends keyof AiModelList, Options extends AiOptions>(
     model: Name,
     inputs: AiModelList[Name]["inputs"],
-    options?: AiOptions,
-  ): Promise<AiModelList[Name]["postProcessedOutputs"]>;
+    options?: Options,
+  ): Promise<
+    Options extends {
+      returnRawResponse: true;
+    }
+      ? Response
+      : AiModelList[Name]["postProcessedOutputs"]
+  >;
   public models(params?: AiModelsSearchParams): Promise<AiModelsSearchObject[]>;
 }
 export type GatewayOptions = {

diff --git a/types/generated-snapshot/2022-03-21/index.d.ts b/types/generated-snapshot/2022-03-21/index.d.ts
@@ -4161,6 +4161,7 @@ interface AiModels {
 }
 type AiOptions = {
   gateway?: GatewayOptions;
+  returnRawResponse?: boolean;
   prefix?: string;
   extraHeaders?: object;
 };
@@ -4195,11 +4196,17 @@ type AiModelListType = Record<string, any>;
 declare abstract class Ai<AiModelList extends AiModelListType = AiModels> {
   aiGatewayLogId: string | null;
   gateway(gatewayId: string): AiGateway;
-  run<Name extends keyof AiModelList>(
+  run<Name extends keyof AiModelList, Options extends AiOptions>(
     model: Name,
     inputs: AiModelList[Name]["inputs"],
-    options?: AiOptions,
-  ): Promise<AiModelList[Name]["postProcessedOutputs"]>;
+    options?: Options,
+  ): Promise<
+    Options extends {
+      returnRawResponse: true;
+    }
+      ? Response
+      : AiModelList[Name]["postProcessedOutputs"]
+  >;
   public models(params?: AiModelsSearchParams): Promise<AiModelsSearchObject[]>;
 }
 type GatewayOptions = {

diff --git a/types/generated-snapshot/2022-03-21/index.ts b/types/generated-snapshot/2022-03-21/index.ts
@@ -4173,6 +4173,7 @@ export interface AiModels {
 }
 export type AiOptions = {
   gateway?: GatewayOptions;
+  returnRawResponse?: boolean;
   prefix?: string;
   extraHeaders?: object;
 };
@@ -4209,11 +4210,17 @@ export declare abstract class Ai<
 > {
   aiGatewayLogId: string | null;
   gateway(gatewayId: string): AiGateway;
-  run<Name extends keyof AiModelList>(
+  run<Name extends keyof AiModelList, Options extends AiOptions>(
     model: Name,
     inputs: AiModelList[Name]["inputs"],
-    options?: AiOptions,
-  ): Promise<AiModelList[Name]["postProcessedOutputs"]>;
+    options?: Options,
+  ): Promise<
+    Options extends {
+      returnRawResponse: true;
+    }
+      ? Response
+      : AiModelList[Name]["postProcessedOutputs"]
+  >;
   public models(params?: AiModelsSearchParams): Promise<AiModelsSearchObject[]>;
 }
 export type GatewayOptions = {

diff --git a/types/generated-snapshot/2022-08-04/index.d.ts b/types/generated-snapshot/2022-08-04/index.d.ts
@@ -4162,6 +4162,7 @@ interface AiModels {
 }
 type AiOptions = {
   gateway?: GatewayOptions;
+  returnRawResponse?: boolean;
   prefix?: string;
   extraHeaders?: object;
 };
@@ -4196,11 +4197,17 @@ type AiModelListType = Record<string, any>;
 declare abstract class Ai<AiModelList extends AiModelListType = AiModels> {
   aiGatewayLogId: string | null;
   gateway(gatewayId: string): AiGateway;
-  run<Name extends keyof AiModelList>(
+  run<Name extends keyof AiModelList, Options extends AiOptions>(
     model: Name,
     inputs: AiModelList[Name]["inputs"],
-    options?: AiOptions,
-  ): Promise<AiModelList[Name]["postProcessedOutputs"]>;
+    options?: Options,
+  ): Promise<
+    Options extends {
+      returnRawResponse: true;
+    }
+      ? Response
+      : AiModelList[Name]["postProcessedOutputs"]
+  >;
   public models(params?: AiModelsSearchParams): Promise<AiModelsSearchObject[]>;
 }
 type GatewayOptions = {

diff --git a/types/generated-snapshot/2022-08-04/index.ts b/types/generated-snapshot/2022-08-04/index.ts
@@ -4174,6 +4174,7 @@ export interface AiModels {
 }
 export type AiOptions = {
   gateway?: GatewayOptions;
+  returnRawResponse?: boolean;
   prefix?: string;
   extraHeaders?: object;
 };
@@ -4210,11 +4211,17 @@ export declare abstract class Ai<
 > {
   aiGatewayLogId: string | null;
   gateway(gatewayId: string): AiGateway;
-  run<Name extends keyof AiModelList>(
+  run<Name extends keyof AiModelList, Options extends AiOptions>(
     model: Name,
     inputs: AiModelList[Name]["inputs"],
-    options?: AiOptions,
-  ): Promise<AiModelList[Name]["postProcessedOutputs"]>;
+    options?: Options,
+  ): Promise<
+    Options extends {
+      returnRawResponse: true;
+    }
+      ? Response
+      : AiModelList[Name]["postProcessedOutputs"]
+  >;
   public models(params?: AiModelsSearchParams): Promise<AiModelsSearchObject[]>;
 }
 export type GatewayOptions = {

diff --git a/types/generated-snapshot/2022-10-31/index.d.ts b/types/generated-snapshot/2022-10-31/index.d.ts
@@ -4165,6 +4165,7 @@ interface AiModels {
 }
 type AiOptions = {
   gateway?: GatewayOptions;
+  returnRawResponse?: boolean;
   prefix?: string;
   extraHeaders?: object;
 };
@@ -4199,11 +4200,17 @@ type AiModelListType = Record<string, any>;
 declare abstract class Ai<AiModelList extends AiModelListType = AiModels> {
   aiGatewayLogId: string | null;
   gateway(gatewayId: string): AiGateway;
-  run<Name extends keyof AiModelList>(
+  run<Name extends keyof AiModelList, Options extends AiOptions>(
     model: Name,
     inputs: AiModelList[Name]["inputs"],
-    options?: AiOptions,
-  ): Promise<AiModelList[Name]["postProcessedOutputs"]>;
+    options?: Options,
+  ): Promise<
+    Options extends {
+      returnRawResponse: true;
+    }
+      ? Response
+      : AiModelList[Name]["postProcessedOutputs"]
+  >;
   public models(params?: AiModelsSearchParams): Promise<AiModelsSearchObject[]>;
 }
 type GatewayOptions = {

diff --git a/types/generated-snapshot/2022-10-31/index.ts b/types/generated-snapshot/2022-10-31/index.ts
@@ -4177,6 +4177,7 @@ export interface AiModels {
 }
 export type AiOptions = {
   gateway?: GatewayOptions;
+  returnRawResponse?: boolean;
   prefix?: string;
   extraHeaders?: object;
 };
@@ -4213,11 +4214,17 @@ export declare abstract class Ai<
 > {
   aiGatewayLogId: string | null;
   gateway(gatewayId: string): AiGateway;
-  run<Name extends keyof AiModelList>(
+  run<Name extends keyof AiModelList, Options extends AiOptions>(
     model: Name,
     inputs: AiModelList[Name]["inputs"],
-    options?: AiOptions,
-  ): Promise<AiModelList[Name]["postProcessedOutputs"]>;
+    options?: Options,
+  ): Promise<
+    Options extends {
+      returnRawResponse: true;
+    }
+      ? Response
+      : AiModelList[Name]["postProcessedOutputs"]
+  >;
   public models(params?: AiModelsSearchParams): Promise<AiModelsSearchObject[]>;
 }
 export type GatewayOptions = {

diff --git a/types/generated-snapshot/2022-11-30/index.d.ts b/types/generated-snapshot/2022-11-30/index.d.ts
@@ -4170,6 +4170,7 @@ interface AiModels {
 }
 type AiOptions = {
   gateway?: GatewayOptions;
+  returnRawResponse?: boolean;
   prefix?: string;
   extraHeaders?: object;
 };
@@ -4204,11 +4205,17 @@ type AiModelListType = Record<string, any>;
 declare abstract class Ai<AiModelList extends AiModelListType = AiModels> {
   aiGatewayLogId: string | null;
   gateway(gatewayId: string): AiGateway;
-  run<Name extends keyof AiModelList>(
+  run<Name extends keyof AiModelList, Options extends AiOptions>(
     model: Name,
     inputs: AiModelList[Name]["inputs"],
-    options?: AiOptions,
-  ): Promise<AiModelList[Name]["postProcessedOutputs"]>;
+    options?: Options,
+  ): Promise<
+    Options extends {
+      returnRawResponse: true;
+    }
+      ? Response
+      : AiModelList[Name]["postProcessedOutputs"]
+  >;
   public models(params?: AiModelsSearchParams): Promise<AiModelsSearchObject[]>;
 }
 type GatewayOptions = {