From 1596e3dc5cefcc22dd29eb7da343c2bfdd1aafc0 Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Fri, 3 Jan 2025 16:58:36 -0500
Subject: [PATCH] chore(js/gemini): added Gemini config metadata

---
 js/plugins/googleai/src/gemini.ts | 20 +++++++++++++++++---
 js/plugins/vertexai/src/gemini.ts | 22 ++++++++++++++++++----
 2 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts
index ebd65adfe..460de3f79 100644
--- a/js/plugins/googleai/src/gemini.ts
+++ b/js/plugins/googleai/src/gemini.ts
@@ -82,14 +82,28 @@ const SafetySettingsSchema = z.object({
 });
 
 export const GeminiConfigSchema = GenerationCommonConfigSchema.extend({
-  safetySettings: z.array(SafetySettingsSchema).optional(),
-  codeExecution: z.union([z.boolean(), z.object({}).strict()]).optional(),
-  contextCache: z.boolean().optional(),
+  version: z.string().describe('Model version').optional(),
+  temperature: z.number().describe('Temperature').min(0).max(2).optional(),
+  maxOutputTokens: z.number().describe('Maximum number of tokens').optional(),
+  topK: z.number().describe('Top-K').min(1).max(40).optional(),
+  topP: z.number().describe('Top-P').min(0).max(1).optional(),
+  stopSequences: z.array(z.string()).describe('Stop sequences').optional(),
+
+  safetySettings: z
+    .array(SafetySettingsSchema)
+    .describe('Safety settings')
+    .optional(),
+  codeExecution: z
+    .union([z.boolean(), z.object({}).strict()])
+    .describe('Enable code execution')
+    .optional(),
+  contextCache: z.boolean().describe('Enable context caching').optional(),
   functionCallingConfig: z
     .object({
       mode: z.enum(['MODE_UNSPECIFIED', 'AUTO', 'ANY', 'NONE']).optional(),
       allowedFunctionNames: z.array(z.string()).optional(),
     })
+    .describe('Function calling config')
     .optional(),
 });
 export type GeminiConfig = z.infer<typeof GeminiConfigSchema>;
diff --git a/js/plugins/vertexai/src/gemini.ts b/js/plugins/vertexai/src/gemini.ts
index 581001619..5c6bf42ca 100644
--- a/js/plugins/vertexai/src/gemini.ts
+++ b/js/plugins/vertexai/src/gemini.ts
@@ -77,10 +77,17 @@ const GoogleSearchRetrievalSchema = z.object({
  * Zod schema of Gemini model options.
  */
 export const GeminiConfigSchema = GenerationCommonConfigSchema.extend({
+  version: z.string().describe('Model version').optional(),
+  temperature: z.number().describe('Temperature').min(0).max(2).optional(),
+  maxOutputTokens: z.number().describe('Maximum number of tokens').optional(),
+  topK: z.number().describe('Top-K').min(1).max(40).optional(),
+  topP: z.number().describe('Top-P').min(0).max(1).optional(),
+  stopSequences: z.array(z.string()).describe('Stop sequences').optional(),
+
   /**
    * GCP region (e.g. us-central1)
    */
-  location: z.string().describe('banana').optional(),
+  location: z.string().describe('GCP region').optional(),
 
   /**
    * Safety filter settings. See: https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/configure-safety-filters#configurable-filters
@@ -110,7 +117,10 @@ export const GeminiConfigSchema = GenerationCommonConfigSchema.extend({
    * }
    * ```
    */
-  safetySettings: z.array(SafetySettingsSchema).optional(),
+  safetySettings: z
+    .array(SafetySettingsSchema)
+    .describe('Safety settings')
+    .optional(),
 
   /**
    * Vertex retrieval options.
@@ -130,7 +140,8 @@ export const GeminiConfigSchema = GenerationCommonConfigSchema.extend({
    *   }
    * ```
    */
-  vertexRetrieval: VertexRetrievalSchema.optional(),
+  vertexRetrieval:
+    VertexRetrievalSchema.describe('Vertex Retrieval').optional(),
 
   /**
    * Google Search retrieval options.
@@ -143,7 +154,9 @@ export const GeminiConfigSchema = GenerationCommonConfigSchema.extend({
    *   }
    * ```
    */
-  googleSearchRetrieval: GoogleSearchRetrievalSchema.optional(),
+  googleSearchRetrieval: GoogleSearchRetrievalSchema.describe(
+    'Google Search Retrieval'
+  ).optional(),
 
   /**
    * Function calling options.
@@ -163,6 +176,7 @@ export const GeminiConfigSchema = GenerationCommonConfigSchema.extend({
       mode: z.enum(['MODE_UNSPECIFIED', 'AUTO', 'ANY', 'NONE']).optional(),
       allowedFunctionNames: z.array(z.string()).optional(),
     })
+    .describe('Function calling config')
     .optional(),
 });