Improving google-cloud and firebase plugin configuration (#886)

- Distinguishing between external options and internal config - Centralizing telemetry defaults and making them environment aware - Exposing configuration methods for use by firebase plugin - Cleaning up more ts in tests
firebase · Sep 11, 2024 · 2f6c7d8 · 2f6c7d8
1 parent 97b7223
commit 2f6c7d8
Show file tree

Hide file tree

Showing 11 changed files with 394 additions and 227 deletions.
diff --git a/js/plugins/firebase/src/index.ts b/js/plugins/firebase/src/index.ts
@@ -18,16 +18,19 @@ import { genkitPlugin, isDevEnv, Plugin } from '@genkit-ai/core';
 import { logger } from '@genkit-ai/core/logging';
 import { FirestoreStateStore } from '@genkit-ai/flow';
 import {
+  configureGcpPlugin,
   GcpLogger,
   GcpOpenTelemetry,
-  TelemetryConfig,
+  GcpTelemetryConfigOptions,
 } from '@genkit-ai/google-cloud';
-import { GoogleAuth } from 'google-auth-library';
+import { JWTInput } from 'google-auth-library';
+import { GcpPluginConfig } from '../../google-cloud/lib/types.js';
 import { FirestoreTraceStore } from './firestoreTraceStore.js';
 export { defineFirestoreRetriever } from './firestoreRetriever.js';
 
-interface FirestorePluginParams {
+export interface FirestorePluginParams {
   projectId?: string;
+  credentials?: JWTInput;
   flowStateStore?: {
     collection?: string;
     databaseId?: string;
@@ -36,46 +39,29 @@ interface FirestorePluginParams {
     collection?: string;
     databaseId?: string;
   };
-  telemetryConfig?: TelemetryConfig;
+  telemetryConfig?: GcpTelemetryConfigOptions;
 }
 
 export const firebase: Plugin<[FirestorePluginParams] | []> = genkitPlugin(
   'firebase',
   async (params?: FirestorePluginParams) => {
-    let authClient;
-    let credentials;
+    const gcpConfig: GcpPluginConfig = await configureGcpPlugin(params);
 
-    // Allow customers to pass in cloud credentials from environment variables
-    // following: https://github.com/googleapis/google-auth-library-nodejs?tab=readme-ov-file#loading-credentials-from-environment-variables
-    if (process.env.GCLOUD_SERVICE_ACCOUNT_CREDS) {
-      const serviceAccountCreds = JSON.parse(
-        process.env.GCLOUD_SERVICE_ACCOUNT_CREDS
+    if (isDevEnv() && !gcpConfig.projectId) {
+      // Helpful warning, since Cloud SDKs probably will not work
+      logger.warn(
+        'WARNING: unable to determine Firebase Project ID. Run "gcloud auth application-default login --project MY_PROJECT_ID"'
       );
-      const authOptions = { credentials: serviceAccountCreds };
-      authClient = new GoogleAuth(authOptions);
-
-      credentials = await authClient.getCredentials();
-    } else {
-      authClient = new GoogleAuth();
     }
 
-    const projectId = params?.projectId || (await authClient.getProjectId());
-
-    const gcpOptions = {
-      projectId,
-      credentials,
-      telemetryConfig: params?.telemetryConfig,
-    };
-
     const flowStateStoreOptions = {
-      projectId,
-      credentials,
+      projectId: gcpConfig.projectId,
+      credentials: gcpConfig.credentials,
       ...params?.flowStateStore,
     };
-
     const traceStoreOptions = {
-      projectId,
-      credentials,
+      projectId: gcpConfig.projectId,
+      credentials: gcpConfig.credentials,
       ...params?.traceStore,
     };
 
@@ -91,26 +77,13 @@ export const firebase: Plugin<[FirestorePluginParams] | []> = genkitPlugin(
       telemetry: {
         instrumentation: {
           id: 'firebase',
-          value: new GcpOpenTelemetry(gcpOptions),
+          value: new GcpOpenTelemetry(gcpConfig),
         },
         logger: {
           id: 'firebase',
-          value: new GcpLogger(gcpOptions),
+          value: new GcpLogger(gcpConfig),
         },
       },
     };
   }
 );
-
-async function getProjectId(authClient: GoogleAuth): Promise<string> {
-  if (isDevEnv()) {
-    return await authClient.getProjectId().catch((err) => {
-      logger.warn(
-        'WARNING: unable to determine Project ID, run "gcloud auth application-default login --project MY_PROJECT_ID"'
-      );
-      return '';
-    });
-  }
-
-  return await authClient.getProjectId();
-}
diff --git a/js/plugins/google-cloud/package.json b/js/plugins/google-cloud/package.json
@@ -60,7 +60,8 @@
     "npm-run-all": "^4.1.5",
     "tsup": "^8.0.2",
     "tsx": "^4.7.0",
-    "typescript": "^4.9.0"
+    "typescript": "^4.9.0",
+    "@genkit-ai/flow": "workspace:*"
   },
   "types": "./lib/index.d.ts",
   "exports": {

diff --git a/js/plugins/google-cloud/src/auth.ts b/js/plugins/google-cloud/src/auth.ts
@@ -0,0 +1,47 @@
+/**
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import { logger } from '@genkit-ai/core/logging';
+import { GoogleAuth } from 'google-auth-library';
+import { GcpPluginOptions } from './types';
+
+/**
+ * Allow customers to pass in cloud credentials from environment variables
+ * following: https://github.com/googleapis/google-auth-library-nodejs?tab=readme-ov-file#loading-credentials-from-environment-variables
+ */
+export async function credentialsFromEnvironment(): Promise<GcpPluginOptions> {
+  let authClient: GoogleAuth;
+  let options: GcpPluginOptions = {};
+
+  if (process.env.GCLOUD_SERVICE_ACCOUNT_CREDS) {
+    const serviceAccountCreds = JSON.parse(
+      process.env.GCLOUD_SERVICE_ACCOUNT_CREDS
+    );
+    const authOptions = { credentials: serviceAccountCreds };
+    authClient = new GoogleAuth(authOptions);
+    options.credentials = await authClient.getCredentials();
+  } else {
+    authClient = new GoogleAuth();
+  }
+  try {
+    const projectId = await authClient.getProjectId();
+    if (projectId && projectId.length > 0) {
+      options.projectId = projectId;
+    }
+  } catch (error) {
+    logger.warn(error);
+  }
+  return options;
+}
diff --git a/js/plugins/google-cloud/src/gcpLogger.ts b/js/plugins/google-cloud/src/gcpLogger.ts
@@ -17,7 +17,7 @@
 import { LoggerConfig } from '@genkit-ai/core';
 import { LoggingWinston } from '@google-cloud/logging-winston';
 import { Writable } from 'stream';
-import { PluginOptions } from './index.js';
+import { GcpPluginConfig } from './types';
 
 /**
  * Additional streams for writing log data to. Useful for unit testing.
@@ -29,11 +29,7 @@ let additionalStream: Writable;
  * logs.
  */
 export class GcpLogger implements LoggerConfig {
-  private readonly options: PluginOptions;
-
-  constructor(options?: PluginOptions) {
-    this.options = options || {};
-  }
+  constructor(private readonly config: GcpPluginConfig) {}
 
   async getLogger(env: string) {
     // Dynamically importing winston here more strictly controls
@@ -54,11 +50,11 @@ export class GcpLogger implements LoggerConfig {
     transports.push(
       this.shouldExport(env)
         ? new LoggingWinston({
-            projectId: this.options.projectId,
+            projectId: this.config.projectId,
             labels: { module: 'genkit' },
             prefix: 'genkit',
             logName: 'genkit_log',
-            credentials: this.options.credentials,
+            credentials: this.config.credentials,
           })
         : new winston.transports.Console()
     );
@@ -73,8 +69,8 @@ export class GcpLogger implements LoggerConfig {
     });
   }
 
-  private shouldExport(env: string) {
-    return this.options.telemetryConfig?.forceDevExport || env !== 'dev';
+  private shouldExport(env?: string) {
+    return this.config.telemetry.export;
   }
 }
 

diff --git a/js/plugins/google-cloud/src/gcpOpenTelemetry.ts b/js/plugins/google-cloud/src/gcpOpenTelemetry.ts
@@ -40,7 +40,6 @@ import {
 } from '@opentelemetry/sdk-metrics';
 import { NodeSDKConfiguration } from '@opentelemetry/sdk-node';
 import {
-  AlwaysOnSampler,
   BatchSpanProcessor,
   InMemorySpanExporter,
   ReadableSpan,
@@ -50,10 +49,10 @@ import {
 import { extractErrorName } from './utils';
 
 import { PathMetadata } from '@genkit-ai/core/tracing';
-import { PluginOptions } from './index.js';
 import { actionTelemetry } from './telemetry/action.js';
 import { flowsTelemetry } from './telemetry/flow.js';
 import { generateTelemetry } from './telemetry/generate.js';
+import { GcpPluginConfig } from './types';
 
 let metricExporter: PushMetricExporter;
 let spanProcessor: BatchSpanProcessor;
@@ -64,52 +63,51 @@ let spanExporter: AdjustingTraceExporter;
  * Metrics, and Logs) to the Google Cloud Operations Suite.
  */
 export class GcpOpenTelemetry implements TelemetryConfig {
-  private readonly options: PluginOptions;
+  private readonly config: GcpPluginConfig;
   private readonly resource: Resource;
 
+  constructor(config: GcpPluginConfig) {
+    this.config = config;
+    this.resource = new Resource({ type: 'global' }).merge(
+      new GcpDetectorSync().detect()
+    );
+  }
+
   /**
    * Log hook for writing trace and span metadata to log messages in the format
    * required by GCP.
    */
   private gcpTraceLogHook = (span: Span, record: any) => {
     const spanContext = span.spanContext();
     const isSampled = !!(spanContext.traceFlags & TraceFlags.SAMPLED);
-    const projectId = this.options.projectId;
+    const projectId = this.config.projectId;
 
     record['logging.googleapis.com/trace'] ??=
       `projects/${projectId}/traces/${spanContext.traceId}`;
     record['logging.googleapis.com/trace_sampled'] ??= isSampled ? '1' : '0';
     record['logging.googleapis.com/spanId'] ??= spanContext.spanId;
   };
 
-  constructor(options?: PluginOptions) {
-    this.options = options || {};
-    this.resource = new Resource({ type: 'global' }).merge(
-      new GcpDetectorSync().detect()
-    );
-  }
-
   getConfig(): Partial<NodeSDKConfiguration> {
     spanProcessor = new BatchSpanProcessor(this.createSpanExporter());
     return {
       resource: this.resource,
       spanProcessor: spanProcessor,
-      sampler: this.options?.telemetryConfig?.sampler || new AlwaysOnSampler(),
+      sampler: this.config.telemetry.sampler,
       instrumentations: this.getInstrumentations(),
       metricReader: this.createMetricReader(),
     };
   }
 
   private createSpanExporter(): SpanExporter {
-    const logIO = this.options.telemetryConfig?.disableLoggingIO ? false : true;
     spanExporter = new AdjustingTraceExporter(
       this.shouldExportTraces()
         ? new TraceExporter({
-            credentials: this.options.credentials,
+            credentials: this.config.credentials,
           })
         : new InMemorySpanExporter(),
-      logIO,
-      this.options.projectId
+      this.config.telemetry.exportIO,
+      this.config.projectId
     );
     return spanExporter;
   }
@@ -120,37 +118,29 @@ export class GcpOpenTelemetry implements TelemetryConfig {
   private createMetricReader(): PeriodicExportingMetricReader {
     metricExporter = this.buildMetricExporter();
     return new PeriodicExportingMetricReader({
-      exportIntervalMillis:
-        this.options?.telemetryConfig?.metricExportIntervalMillis || 300_000,
-      exportTimeoutMillis:
-        this.options?.telemetryConfig?.metricExportTimeoutMillis || 300_000,
+      exportIntervalMillis: this.config.telemetry.metricExportIntervalMillis,
+      exportTimeoutMillis: this.config.telemetry.metricExportTimeoutMillis,
       exporter: metricExporter,
     });
   }
 
   /** Gets all open telemetry instrumentations as configured by the plugin. */
   private getInstrumentations() {
-    if (this.options?.telemetryConfig?.autoInstrumentation) {
+    if (this.config.telemetry.autoInstrumentation) {
       return getNodeAutoInstrumentations(
-        this.options?.telemetryConfig?.autoInstrumentationConfig || {}
+        this.config.telemetry.autoInstrumentationConfig
       ).concat(this.getDefaultLoggingInstrumentations());
     }
     return this.getDefaultLoggingInstrumentations();
   }
 
   private shouldExportTraces(): boolean {
-    return (
-      (this.options.telemetryConfig?.forceDevExport ||
-        process.env.GENKIT_ENV !== 'dev') &&
-      !this.options.telemetryConfig?.disableTraces
-    );
+    return this.config.telemetry.export && !this.config.telemetry.disableTraces;
   }
 
   private shouldExportMetrics(): boolean {
     return (
-      (this.options.telemetryConfig?.forceDevExport ||
-        process.env.GENKIT_ENV !== 'dev') &&
-      !this.options.telemetryConfig?.disableMetrics
+      this.config.telemetry.export && !this.config.telemetry.disableMetrics
     );
   }
 
@@ -165,12 +155,12 @@ export class GcpOpenTelemetry implements TelemetryConfig {
   private buildMetricExporter(): PushMetricExporter {
     const exporter: PushMetricExporter = this.shouldExportMetrics()
       ? new MetricExporter({
-          projectId: this.options.projectId,
+          projectId: this.config.projectId,
           userAgent: {
             product: 'genkit',
             version: GENKIT_VERSION,
           },
-          credentials: this.options.credentials,
+          credentials: this.config.credentials,
         })
       : new InMemoryMetricExporter(AggregationTemporality.DELTA);
     exporter.selectAggregation = (instrumentType: InstrumentType) => {