Initial prototype

Author: dlarocque

common/api-review/vertexai.api.md

@@ -324,7 +324,7 @@ export interface GenerativeContentBlob {
 
 // @public
 export class GenerativeModel extends VertexAIModel {
-    constructor(vertexAI: VertexAI, modelParams: ModelParams, requestOptions?: RequestOptions);
+    constructor(vertexAI: VertexAI, modelParams: ModelParams, requestOptions?: RequestOptions, developerAPIEnabled?: boolean);
     countTokens(request: CountTokensRequest | string | Array<string | Part>): Promise<CountTokensResponse>;
     generateContent(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentResult>;
     generateContentStream(request: GenerateContentRequest | string | Array<string | Part>): Promise<GenerateContentStreamResult>;
@@ -344,7 +344,7 @@ export class GenerativeModel extends VertexAIModel {
 }
 
 // @public
-export function getGenerativeModel(vertexAI: VertexAI, modelParams: ModelParams, requestOptions?: RequestOptions): GenerativeModel;
+export function getGenerativeModel(vertexAI: VertexAI, modelParams: ModelParams, requestOptions?: RequestOptions, enableDeveloperAPI?: boolean): GenerativeModel;
 
 // @beta
 export function getImagenModel(vertexAI: VertexAI, modelParams: ImagenModelParams, requestOptions?: RequestOptions): ImagenModel;
@@ -776,6 +776,8 @@ export interface UsageMetadata {
 export interface VertexAI {
     app: FirebaseApp;
     // (undocumented)
+    developerAPIEnabled: boolean;
+    // (undocumented)
     location: string;
 }
 
@@ -806,15 +808,17 @@ export const enum VertexAIErrorCode {
 // @public
 export abstract class VertexAIModel {
     // @internal
-    protected constructor(vertexAI: VertexAI, modelName: string);
+    protected constructor(vertexAI: VertexAI, modelName: string, developerAPIEnabled?: boolean);
     // @internal (undocumented)
     protected _apiSettings: ApiSettings;
     readonly model: string;
-    static normalizeModelName(modelName: string): string;
+    static normalizeModelName(modelName: string, developerAPIEnabled?: boolean): string;
 }
 
 // @public
 export interface VertexAIOptions {
+    // (undocumented)
+    developerAPIEnabled: boolean;
     // (undocumented)
     location?: string;
 }

packages/vertexai/src/api.ts

@@ -29,6 +29,7 @@ import {
 } from './types';
 import { VertexAIError } from './errors';
 import { VertexAIModel, GenerativeModel, ImagenModel } from './models';
+import { createInstanceIdentifier } from './helpers';
 
 export { ChatSession } from './methods/chat-session';
 export * from './requests/schema-builder';
@@ -57,8 +58,9 @@ export function getVertexAI(
   // Dependencies
   const vertexProvider: Provider<'vertexAI'> = _getProvider(app, VERTEX_TYPE);
 
+  const identifier = createInstanceIdentifier(options?.developerAPIEnabled, options?.location);
   return vertexProvider.getImmediate({
-    identifier: options?.location || DEFAULT_LOCATION
+    identifier
   });
 }
 
@@ -71,15 +73,16 @@ export function getVertexAI(
 export function getGenerativeModel(
   vertexAI: VertexAI,
   modelParams: ModelParams,
-  requestOptions?: RequestOptions
+  requestOptions?: RequestOptions,
+  enableDeveloperAPI?: boolean
 ): GenerativeModel {
   if (!modelParams.model) {
     throw new VertexAIError(
       VertexAIErrorCode.NO_MODEL,
       `Must provide a model name. Example: getGenerativeModel({ model: 'my-model-name' })`
     );
   }
-  return new GenerativeModel(vertexAI, modelParams, requestOptions);
+  return new GenerativeModel(vertexAI, modelParams, requestOptions, enableDeveloperAPI);
 }
 
 /**

packages/vertexai/src/constants.ts

@@ -23,6 +23,8 @@ export const DEFAULT_LOCATION = 'us-central1';
 
 export const DEFAULT_BASE_URL = 'https://firebasevertexai.googleapis.com';
 
+export const DEVELOPER_API_BASE_URL = "https://generativelanguage.googleapis.com";
+
 export const DEFAULT_API_VERSION = 'v1beta';
 
 export const PACKAGE_VERSION = version;

packages/vertexai/src/helpers.ts

@@ -0,0 +1,31 @@
+import { DEFAULT_LOCATION } from "./constants";
+
+/**
+ * @internal
+ */
+export function createInstanceIdentifier(developerAPIEnabled?: boolean, location?: string): string {
+  if (developerAPIEnabled) {
+    return 'developerAPI';
+  } else {
+    return `vertexAI/${location || DEFAULT_LOCATION}`;
+  }
+}
+
+/**
+ * @internal
+ */
+export function parseInstanceIdentifier(instanceIdentifier: string): { developerAPIEnabled: boolean, location?: string } {
+  const identifierParts = instanceIdentifier.split("/");
+  if (identifierParts[0] === 'developerAPI') {
+    return {
+      developerAPIEnabled: true,
+      location: undefined
+    }
+  } else {
+    const location = identifierParts[1];
+    return {
+      developerAPIEnabled: false,
+      location
+    }
+  }
+}

packages/vertexai/src/index.ts

@@ -24,8 +24,10 @@
 import { registerVersion, _registerComponent } from '@firebase/app';
 import { VertexAIService } from './service';
 import { VERTEX_TYPE } from './constants';
-import { Component, ComponentType } from '@firebase/component';
+import { Component, ComponentType, InstanceFactoryOptions } from '@firebase/component';
 import { name, version } from '../package.json';
+import { parseInstanceIdentifier } from './helpers';
+import { VertexAIOptions } from './public-types';
 
 declare global {
   interface Window {
@@ -37,12 +39,23 @@ function registerVertex(): void {
   _registerComponent(
     new Component(
       VERTEX_TYPE,
-      (container, { instanceIdentifier: location }) => {
+      (container, options) => {
         // getImmediate for FirebaseApp will always succeed
         const app = container.getProvider('app').getImmediate();
         const auth = container.getProvider('auth-internal');
         const appCheckProvider = container.getProvider('app-check-internal');
-        return new VertexAIService(app, auth, appCheckProvider, { location });
+
+        let vertexAIOptions: VertexAIOptions;
+        if (options.instanceIdentifier) {
+          vertexAIOptions = parseInstanceIdentifier(options.instanceIdentifier);
+        } else {
+          vertexAIOptions = {
+            developerAPIEnabled: false,
+            location: undefined
+          }
+        }
+
+        return new VertexAIService(app, auth, appCheckProvider, vertexAIOptions);
       },
       ComponentType.PUBLIC
     ).setMultipleInstances(true)

packages/vertexai/src/methods/generate-content.ts

@@ -56,7 +56,7 @@ export async function generateContent(
     apiSettings,
     /* stream */ false,
     JSON.stringify(params),
-    requestOptions
+    requestOptions,
   );
   const responseJson: GenerateContentResponse = await response.json();
   const enhancedResponse = createEnhancedContentResponse(responseJson);

packages/vertexai/src/models/generative-model.ts

@@ -59,7 +59,7 @@ export class GenerativeModel extends VertexAIModel {
   constructor(
     vertexAI: VertexAI,
     modelParams: ModelParams,
-    requestOptions?: RequestOptions
+    requestOptions?: RequestOptions,
   ) {
     super(vertexAI, modelParams.model);
     this.generationConfig = modelParams.generationConfig || {};

packages/vertexai/src/models/vertexai-model.ts

@@ -56,8 +56,6 @@ export abstract class VertexAIModel {
    * @internal
    */
   protected constructor(vertexAI: VertexAI, modelName: string) {
-    this.model = VertexAIModel.normalizeModelName(modelName);
-
     if (!vertexAI.app?.options?.apiKey) {
       throw new VertexAIError(
         VertexAIErrorCode.NO_API_KEY,
@@ -72,7 +70,8 @@ export abstract class VertexAIModel {
       this._apiSettings = {
         apiKey: vertexAI.app.options.apiKey,
         project: vertexAI.app.options.projectId,
-        location: vertexAI.location
+        developerAPIEnabled: vertexAI.developerAPIEnabled,
+        location: vertexAI.location,
       };
 
       if (
@@ -92,6 +91,8 @@ export abstract class VertexAIModel {
         this._apiSettings.getAuthToken = () =>
           (vertexAI as VertexAIService).auth!.getToken();
       }
+
+      this.model = this.normalizeModelName(modelName);
     }
   }
 
@@ -101,19 +102,23 @@ export abstract class VertexAIModel {
    * @param modelName - The model name to normalize.
    * @returns The fully qualified model resource name.
    */
-  static normalizeModelName(modelName: string): string {
+  normalizeModelName(modelName: string): string {
     let model: string;
-    if (modelName.includes('/')) {
-      if (modelName.startsWith('models/')) {
-        // Add 'publishers/google' if the user is only passing in 'models/model-name'.
-        model = `publishers/google/${modelName}`;
+    if (this._apiSettings.developerAPIEnabled) {
+      model = `models/${modelName}`;
+    } else {
+      if (modelName.includes('/')) {
+        if (modelName.startsWith('models/')) {
+          // Add 'publishers/google' if the user is only passing in 'models/model-name'.
+          model = `publishers/google/${modelName}`;
+        } else {
+          // Any other custom format (e.g. tuned models) must be passed in correctly.
+          model = modelName;
+        }
       } else {
-        // Any other custom format (e.g. tuned models) must be passed in correctly.
-        model = modelName;
+        // If path is not included, assume it's a non-tuned model.
+        model = `publishers/google/models/${modelName}`;
       }
-    } else {
-      // If path is not included, assume it's a non-tuned model.
-      model = `publishers/google/models/${modelName}`;
     }
 
     return model;

packages/vertexai/src/public-types.ts

@@ -28,13 +28,15 @@ export interface VertexAI {
    * The {@link @firebase/app#FirebaseApp} this <code>{@link VertexAI}</code> instance is associated with.
    */
   app: FirebaseApp;
-  location: string;
+  developerAPIEnabled: boolean;
+  location: string; // This is only applicable if we're using the VertexAI API.
 }
 
 /**
  * Options when initializing the Vertex AI in Firebase SDK.
  * @public
  */
 export interface VertexAIOptions {
+  developerAPIEnabled: boolean;
   location?: string;
 }

packages/vertexai/src/requests/request.ts

@@ -22,6 +22,7 @@ import {
   DEFAULT_API_VERSION,
   DEFAULT_BASE_URL,
   DEFAULT_FETCH_TIMEOUT_MS,
+  DEVELOPER_API_BASE_URL,
   LANGUAGE_TAG,
   PACKAGE_VERSION
 } from '../constants';
@@ -40,17 +41,23 @@ export class RequestUrl {
     public task: Task,
     public apiSettings: ApiSettings,
     public stream: boolean,
-    public requestOptions?: RequestOptions
+    public requestOptions?: RequestOptions,
   ) {}
   toString(): string {
     // TODO: allow user-set option if that feature becomes available
     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}`;
+    let url;
+    if (this.apiSettings.developerAPIEnabled) {
+      const baseUrl = this.requestOptions?.baseUrl || DEVELOPER_API_BASE_URL;
+      url = `${baseUrl}/${apiVersion}/${this.model}:${this.task}`;
+    } else {
+      const baseUrl = this.requestOptions?.baseUrl || DEFAULT_BASE_URL;
+      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';
     }