refactor: remove google defaults

danny-avila · Feb 5, 2025 · de6de73 · de6de73
1 parent 0312d4f
commit de6de73
Showing 1 changed file with 21 additions and 5 deletions.
diff --git a/api/app/clients/GoogleClient.js b/api/app/clients/GoogleClient.js
@@ -566,7 +566,15 @@ class GoogleClient extends BaseClient {
 
     if (this.project_id != null) {
       logger.debug('Creating VertexAI client');
-      return new ChatVertexAI(clientOptions);
+      clientOptions.streaming = true;
+      const client = new ChatVertexAI(clientOptions);
+      client.temperature = clientOptions.temperature;
+      client.topP = clientOptions.topP;
+      client.topK = clientOptions.topK;
+      client.topLogprobs = clientOptions.topLogprobs;
+      client.frequencyPenalty = clientOptions.frequencyPenalty;
+      client.presencePenalty = clientOptions.presencePenalty;
+      client.maxOutputTokens = clientOptions.maxOutputTokens;
     } else if (!EXCLUDED_GENAI_MODELS.test(model)) {
       logger.debug('Creating GenAI client');
       return new GenAI(this.apiKey).getGenerativeModel({ model }, requestOptions);
@@ -577,7 +585,7 @@ class GoogleClient extends BaseClient {
   }
 
   initializeClient() {
-    let clientOptions = { ...this.modelOptions, maxRetries: 2 };
+    let clientOptions = { ...this.modelOptions };
 
     if (this.project_id) {
       clientOptions['authOptions'] = {
@@ -681,9 +689,17 @@ class GoogleClient extends BaseClient {
       }
 
       for await (const chunk of stream) {
-        usageMetadata = !usageMetadata
-          ? chunk?.usage_metadata
-          : concat(usageMetadata, chunk?.usage_metadata);
+        if (chunk?.usage_metadata) {
+          const metadata = chunk.usage_metadata;
+          for (const key in metadata) {
+            if (Number.isNaN(metadata[key])) {
+              delete metadata[key];
+            }
+          }
+
+          usageMetadata = !usageMetadata ? metadata : concat(usageMetadata, metadata);
+        }
+
         const chunkText = chunk?.content ?? chunk;
         await this.generateTextStream(chunkText, onProgress, {
           delay,