request.ts

Author: russellwheatley

packages/ai/lib/requests/request.ts

@@ -15,8 +15,8 @@
  * limitations under the License.
  */
 import { Platform } from 'react-native';
-import { ErrorDetails, RequestOptions, VertexAIErrorCode } from '../types';
-import { VertexAIError } from '../errors';
+import { AIErrorCode, ErrorDetails, RequestOptions } from '../types';
+import { AIError } from '../errors';
 import { ApiSettings } from '../types/internal';
 import {
   DEFAULT_API_VERSION,
@@ -26,11 +26,13 @@ import {
   PACKAGE_VERSION,
 } from '../constants';
 import { logger } from '../logger';
+import { GoogleAIBackend, VertexAIBackend } from '../backend';
 
 export enum Task {
   GENERATE_CONTENT = 'generateContent',
   STREAM_GENERATE_CONTENT = 'streamGenerateContent',
   COUNT_TOKENS = 'countTokens',
+  PREDICT = 'predict',
 }
 
 export class RequestUrl {
@@ -58,29 +60,40 @@ export class RequestUrl {
       }
       return emulatorUrl;
     }
+    const url = new URL(this.baseUrl); // Throws if the URL is invalid
+    url.pathname = `/${this.apiVersion}/${this.modelPath}:${this.task}`;
+    url.search = this.queryParams.toString();
+    return url.toString();
+  }
 
-    const apiVersion = DEFAULT_API_VERSION;
-    const baseUrl = this.requestOptions?.baseUrl || DEFAULT_BASE_URL;
-    let url = `${baseUrl}/${apiVersion}`;
-    url += `/projects/${this.apiSettings.project}`;
-    url += `/locations/${this.apiSettings.location}`;
-    url += `/${this.model}`;
-    url += `:${this.task}`;
-    if (this.stream) {
-      url += '?alt=sse';
+  private get baseUrl(): string {
+    return this.requestOptions?.baseUrl || DEFAULT_BASE_URL;
+  }
+
+  private get apiVersion(): string {
+    return DEFAULT_API_VERSION; // TODO: allow user-set options if that feature becomes available
+  }
+
+  private get modelPath(): string {
+    if (this.apiSettings.backend instanceof GoogleAIBackend) {
+      return `projects/${this.apiSettings.project}/${this.model}`;
+    } else if (this.apiSettings.backend instanceof VertexAIBackend) {
+      return `projects/${this.apiSettings.project}/locations/${this.apiSettings.backend.location}/${this.model}`;
+    } else {
+      throw new AIError(
+        AIErrorCode.ERROR,
+        `Invalid backend: ${JSON.stringify(this.apiSettings.backend)}`,
+      );
     }
-    return url;
   }
 
-  /**
-   * If the model needs to be passed to the backend, it needs to
-   * include project and location path.
-   */
-  get fullModelString(): string {
-    let modelString = `projects/${this.apiSettings.project}`;
-    modelString += `/locations/${this.apiSettings.location}`;
-    modelString += `/${this.model}`;
-    return modelString;
+  private get queryParams(): URLSearchParams {
+    const params = new URLSearchParams();
+    if (this.stream) {
+      params.set('alt', 'sse');
+    }
+
+    return params;
   }
 }
 
@@ -99,6 +112,9 @@ export async function getHeaders(url: RequestUrl): Promise<Headers> {
   headers.append('Content-Type', 'application/json');
   headers.append('x-goog-api-client', getClientHeaders());
   headers.append('x-goog-api-key', url.apiSettings.apiKey);
+  if (url.apiSettings.automaticDataCollectionEnabled) {
+    headers.append('X-Firebase-Appid', url.apiSettings.appId);
+  }
   if (url.apiSettings.getAppCheckToken) {
     let appCheckToken;
 
@@ -154,6 +170,7 @@ export async function makeRequest(
   let fetchTimeoutId: string | number | NodeJS.Timeout | undefined;
   try {
     const request = await constructRequest(model, task, apiSettings, stream, body, requestOptions);
+    // Timeout is 180s by default
     const timeoutMillis =
       requestOptions?.timeout != null && requestOptions.timeout >= 0
         ? requestOptions.timeout
@@ -192,9 +209,9 @@ export async function makeRequest(
           ),
         )
       ) {
-        throw new VertexAIError(
-          VertexAIErrorCode.API_NOT_ENABLED,
-          `The Vertex AI in Firebase SDK requires the Vertex AI in Firebase ` +
+        throw new AIError(
+          AIErrorCode.API_NOT_ENABLED,
+          `The Firebase AI SDK requires the Firebase AI ` +
             `API ('firebasevertexai.googleapis.com') to be enabled in your ` +
             `Firebase project. Enable this API by visiting the Firebase Console ` +
             `at https://console.firebase.google.com/project/${url.apiSettings.project}/genai/ ` +
@@ -208,8 +225,8 @@ export async function makeRequest(
           },
         );
       }
-      throw new VertexAIError(
-        VertexAIErrorCode.FETCH_ERROR,
+      throw new AIError(
+        AIErrorCode.FETCH_ERROR,
         `Error fetching from ${url}: [${response.status} ${response.statusText}] ${message}`,
         {
           status: response.status,
@@ -221,14 +238,11 @@ export async function makeRequest(
   } catch (e) {
     let err = e as Error;
     if (
-      (e as VertexAIError).code !== VertexAIErrorCode.FETCH_ERROR &&
-      (e as VertexAIError).code !== VertexAIErrorCode.API_NOT_ENABLED &&
+      (e as AIError).code !== AIErrorCode.FETCH_ERROR &&
+      (e as AIError).code !== AIErrorCode.API_NOT_ENABLED &&
       e instanceof Error
     ) {
-      err = new VertexAIError(
-        VertexAIErrorCode.ERROR,
-        `Error fetching from ${url.toString()}: ${e.message}`,
-      );
+      err = new AIError(AIErrorCode.ERROR, `Error fetching from ${url.toString()}: ${e.message}`);
       err.stack = e.stack;
     }