Define HybridParams

Author: erikeldridge

common/api-review/vertexai.api.md

@@ -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, onCloudOrHybridParams: ModelParams | HybridParams, requestOptions?: RequestOptions): GenerativeModel;
 
 // @beta
 export function getImagenModel(vertexAI: VertexAI, modelParams: ImagenModelParams, requestOptions?: RequestOptions): ImagenModel;
@@ -500,6 +500,28 @@ export interface ImagenSafetySettings {
     safetyFilterLevel?: ImagenSafetyFilterLevel;
 }
 
+// @public
+export interface HybridParams {
+    // (undocumented)
+    mode?: InferenceMode;
+    // (undocumented)
+    onCloudParams?: ModelParams;
+    // Warning: (ae-forgotten-export) The symbol "AILanguageModelCreateOptionsWithSystemPrompt" needs to be exported by the entry point index.d.ts
+    //
+    // (undocumented)
+    onDeviceParams?: AILanguageModelCreateOptionsWithSystemPrompt;
+}
+
+// @public
+export enum InferenceMode {
+    // (undocumented)
+    ONLY_ON_CLOUD = "ONLY_ON_CLOUD",
+    // (undocumented)
+    ONLY_ON_DEVICE = "ONLY_ON_DEVICE",
+    // (undocumented)
+    PREFER_ON_DEVICE = "PREFER_ON_DEVICE"
+}
+
 // @public
 export interface InlineDataPart {
     // (undocumented)

packages/vertexai/src/api.ts

@@ -23,6 +23,7 @@ import { VertexAIService } from './service';
 import { VertexAI, VertexAIOptions } from './public-types';
 import {
   ImagenModelParams,
+  HybridParams,
   ModelParams,
   RequestOptions,
   VertexAIErrorCode
@@ -70,16 +71,27 @@ export function getVertexAI(
  */
 export function getGenerativeModel(
   vertexAI: VertexAI,
-  modelParams: ModelParams,
+  onCloudOrHybridParams: ModelParams | HybridParams,
   requestOptions?: RequestOptions
 ): GenerativeModel {
-  if (!modelParams.model) {
+  // Disambiguates onCloudOrHybridParams input.
+  const hybridParams = onCloudOrHybridParams as HybridParams;
+  let onCloudParams: ModelParams;
+  if (hybridParams.mode) {
+    onCloudParams = hybridParams.onCloudParams || {
+      model: 'gemini-2.0-flash-lite'
+    };
+  } else {
+    onCloudParams = onCloudOrHybridParams as ModelParams;
+  }
+
+  if (!onCloudParams.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, onCloudParams, requestOptions);
 }
 
 /**

packages/vertexai/src/types/enums.ts

@@ -240,3 +240,13 @@ export enum Modality {
    */
   DOCUMENT = 'DOCUMENT'
 }
+
+/**
+ * Determines whether inference happens on-device or on-cloud.
+ * @public
+ */
+export enum InferenceMode {
+  PREFER_ON_DEVICE = 'PREFER_ON_DEVICE',
+  ONLY_ON_DEVICE = 'ONLY_ON_DEVICE',
+  ONLY_ON_CLOUD = 'ONLY_ON_CLOUD'
+}

packages/vertexai/src/types/language-model.ts

@@ -38,12 +38,12 @@ enum Availability {
   'downloading',
   'available'
 }
-interface LanguageModelCreateCoreOptions {
+export interface LanguageModelCreateCoreOptions {
   topK?: number;
   temperature?: number;
   expectedInputs?: LanguageModelExpectedInput[];
 }
-interface LanguageModelCreateOptions extends LanguageModelCreateCoreOptions {
+export interface LanguageModelCreateOptions extends LanguageModelCreateCoreOptions {
   signal?: AbortSignal;
   systemPrompt?: string;
   initialPrompts?: LanguageModelInitialPrompts;

packages/vertexai/src/types/requests.ts

@@ -17,11 +17,13 @@
 
 import { TypedSchema } from '../requests/schema-builder';
 import { Content, Part } from './content';
+import {LanguageModelCreateOptions} from './language-model';
 import {
   FunctionCallingMode,
   HarmBlockMethod,
   HarmBlockThreshold,
-  HarmCategory
+  HarmCategory,
+  InferenceMode
 } from './enums';
 import { ObjectSchemaInterface, SchemaRequest } from './schema';
 
@@ -213,3 +215,13 @@ export interface FunctionCallingConfig {
   mode?: FunctionCallingMode;
   allowedFunctionNames?: string[];
 }
+
+/**
+ * Configures on-device and on-cloud inference.
+ * @public
+ */
+export interface HybridParams {
+  mode?: InferenceMode;
+  onDeviceParams?: LanguageModelCreateOptions;
+  onCloudParams?: ModelParams;
+}