[JS] Moved Firestore implementation from Firebase plugin to Google Cloud plugin.

js/ai/src/generate.ts

@@ -16,19 +16,19 @@
 
 import {
   Action,
-  config as genkitConfig,
   GenkitError,
-  runWithStreamingCallback,
   StreamingCallback,
+  config as genkitConfig,
+  runWithStreamingCallback,
 } from '@genkit-ai/core';
 import { lookupAction } from '@genkit-ai/core/registry';
 import { toJsonSchema, validateSchema } from '@genkit-ai/core/schema';
 import { z } from 'zod';
 import { DocumentData } from './document.js';
 import { extractJson } from './extract.js';
 import {
-  generateAction,
   GenerateUtilParamSchema,
+  generateAction,
   inferRoleFromParts,
 } from './generateAction.js';
 import {
@@ -47,7 +47,7 @@ import {
   ToolRequestPart,
   ToolResponsePart,
 } from './model.js';
-import { resolveTools, ToolArgument, toToolDefinition } from './tool.js';
+import { ToolArgument, resolveTools, toToolDefinition } from './tool.js';
 
 /**
  * Message represents a single role's contribution to a generation. Each message

js/plugins/firebase/package.json

@@ -37,7 +37,6 @@
     "zod": "^3.22.4"
   },
   "peerDependencies": {
-    "@google-cloud/firestore": "^7.6.0",
     "firebase-admin": "^12.2.0",
     "firebase-functions": "^4.8.0 || ^5.0.0",
     "@genkit-ai/ai": "workspace:*",

js/plugins/firebase/src/auth.ts

@@ -21,20 +21,35 @@ import * as z from 'zod';
 import { FunctionFlowAuth } from './functions.js';
 import { initializeAppIfNecessary } from './helpers.js';
 
+/**
+ * Provides a Firebase Auth implementation for Genkit with required auth for direct calls.
+ */
 export function firebaseAuth<I extends z.ZodTypeAny>(
   policy: (user: DecodedIdToken, input: z.infer<I>) => void | Promise<void>
 ): FunctionFlowAuth<I>;
+
+/**
+ * Provides a Firebase Auth implementation for Genkit with required auth for direct calls.
+ */
 export function firebaseAuth<I extends z.ZodTypeAny>(
   policy: (user: DecodedIdToken, input: z.infer<I>) => void | Promise<void>,
   config: { required: true }
 ): FunctionFlowAuth<I>;
+
+/**
+ * Provides a Firebase Auth implementation for Genkit with optional auth for direct calls.
+ */
 export function firebaseAuth<I extends z.ZodTypeAny>(
   policy: (
     user: DecodedIdToken | undefined,
     input: z.infer<I>
   ) => void | Promise<void>,
   config: { required: false }
 ): FunctionFlowAuth<I>;
+
+/**
+ * Provides a Firebase Auth implementation for Genkit.
+ */
 export function firebaseAuth<I extends z.ZodTypeAny>(
   policy: (user: DecodedIdToken, input: z.infer<I>) => void | Promise<void>,
   config?: { required: boolean }

js/plugins/firebase/src/functions.ts

@@ -81,7 +81,7 @@ export function onFlow<
       invoker: async (flow, data, streamingCallback) => {
         const responseJson = await callHttpsFunction(
           flow.name,
-          await getLocation(),
+          getLocation(),
           data,
           streamingCallback
         );
@@ -124,7 +124,7 @@ function wrapHttpsFlow<
     async (req, res) => {
       if (config.enforceAppCheck) {
         if (
-          !(await appCheckValid(
+          !(await isAppCheckValid(
             req.headers['x-firebase-appcheck'],
             !!config.consumeAppCheckToken
           ))
@@ -152,16 +152,22 @@ function wrapHttpsFlow<
   );
 }
 
-async function appCheckValid(
-  tok: string | string[] | undefined,
+/**
+ * Checks if the App Check token is valid.
+ */
+async function isAppCheckValid(
+  token: string | string[] | undefined,
   consume: boolean
 ): Promise<boolean> {
-  if (typeof tok !== 'string') return false;
+  if (typeof token !== 'string') {
+    return false;
+  }
   initializeAppIfNecessary();
   try {
-    const appCheckClaims = await getAppCheck().verifyToken(tok, { consume });
-
-    if (consume && appCheckClaims.alreadyConsumed) return false;
+    const appCheckClaims = await getAppCheck().verifyToken(token, { consume });
+    if (consume && appCheckClaims.alreadyConsumed) {
+      return false;
+    }
     return true;
   } catch (e) {
     return false;

js/plugins/firebase/src/index.ts

@@ -18,30 +18,48 @@ import { genkitPlugin, isDevEnv, Plugin } from '@genkit-ai/core';
 import { logger } from '@genkit-ai/core/logging';
 import { FirestoreStateStore } from '@genkit-ai/flow';
 import {
+  Credentials,
+  FirestoreTraceStore,
   GcpLogger,
   GcpOpenTelemetry,
   TelemetryConfig,
 } from '@genkit-ai/google-cloud';
 import { GoogleAuth } from 'google-auth-library';
-import { FirestoreTraceStore } from './firestoreTraceStore.js';
-export { defineFirestoreRetriever } from './firestoreRetriever.js';
 
-interface FirestorePluginParams {
+export { defineFirestoreRetriever } from '@genkit-ai/google-cloud';
+
+/**
+ * Parameters for the Firebase plugin.
+ */
+interface FirebasePluginParams {
+  /** Firebase project ID. */
   projectId?: string;
+  /** Configuration for the Firestore-based flow state store. */
   flowStateStore?: {
+    /** Firestore collection to use. If not provided, the default collection is used. */
     collection?: string;
+    /** Firestore database ID to use. If not provided, the default database ID is used. */
     databaseId?: string;
   };
+  /** Configuration for the Firestore-based trace store. */
   traceStore?: {
+    /** Firestore collection to use. If not provided, the default collection is used. */
     collection?: string;
+    /** Firestore database ID to use. If not provided, the default database ID is used. */
     databaseId?: string;
   };
+  /** Configuration for the OpenTelemetry telemetry exporter. */
   telemetryConfig?: TelemetryConfig;
+  /** Credentials to use for the Google Cloud API. */
+  credentials?: Credentials;
 }
 
-export const firebase: Plugin<[FirestorePluginParams] | []> = genkitPlugin(
+/**
+ * Provides a Firebase plugin for Genkit.
+ */
+export const firebase: Plugin<[FirebasePluginParams] | []> = genkitPlugin(
   'firebase',
-  async (params?: FirestorePluginParams) => {
+  async (params?: FirebasePluginParams) => {
     let authClient;
     let credentials;
 
@@ -59,9 +77,9 @@ export const firebase: Plugin<[FirestorePluginParams] | []> = genkitPlugin(
       authClient = new GoogleAuth();
     }
 
-    const projectId = params?.projectId || (await authClient.getProjectId());
+    const projectId = params?.projectId || (await getProjectId(authClient));
 
-    const gcpOptions = {
+    const paramsWithProjectIdAndCreds = {
       projectId,
       credentials,
       telemetryConfig: params?.telemetryConfig,
@@ -91,11 +109,11 @@ export const firebase: Plugin<[FirestorePluginParams] | []> = genkitPlugin(
       telemetry: {
         instrumentation: {
           id: 'firebase',
-          value: new GcpOpenTelemetry(gcpOptions),
+          value: new GcpOpenTelemetry(paramsWithProjectIdAndCreds),
         },
         logger: {
           id: 'firebase',
-          value: new GcpLogger(gcpOptions),
+          value: new GcpLogger(paramsWithProjectIdAndCreds),
         },
       },
     };
@@ -114,3 +132,5 @@ async function getProjectId(authClient: GoogleAuth): Promise<string> {
 
   return await authClient.getProjectId();
 }
+
+export default firebase;

js/plugins/google-cloud/package.json

@@ -53,7 +53,9 @@
   },
   "peerDependencies": {
     "@genkit-ai/ai": "workspace:*",
-    "@genkit-ai/core": "workspace:*"
+    "@genkit-ai/core": "workspace:*",
+    "@genkit-ai/flow": "workspace:*",
+    "@google-cloud/firestore": "^7.6.0"
   },
   "devDependencies": {
     "@types/node": "^20.11.16",

js/plugins/firebase/src/firestoreTraceStore.ts → js/plugins/google-cloud/src/firestoreTraceStore.ts

@@ -30,7 +30,7 @@ import { randomUUID } from 'crypto';
 const DOC_MAX_SIZE = 1_000_000;
 
 /** Allow customers to set service account credentials via an environment variable. */
-interface Credentials {
+export interface Credentials {
   client_email?: string;
   private_key?: string;
 }

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 { GoogleCloudPluginParams } from './index.js';
 
 /**
  * Additional streams for writing log data to. Useful for unit testing.
@@ -29,10 +29,10 @@ let additionalStream: Writable;
  * logs.
  */
 export class GcpLogger implements LoggerConfig {
-  private readonly options: PluginOptions;
+  private readonly params: GoogleCloudPluginParams;
 
-  constructor(options?: PluginOptions) {
-    this.options = options || {};
+  constructor(params?: GoogleCloudPluginParams) {
+    this.params = params || {};
   }
 
   async getLogger(env: string) {
@@ -54,11 +54,11 @@ export class GcpLogger implements LoggerConfig {
     transports.push(
       this.shouldExport(env)
         ? new LoggingWinston({
-            projectId: this.options.projectId,
+            projectId: this.params.projectId,
             labels: { module: 'genkit' },
             prefix: 'genkit',
             logName: 'genkit_log',
-            credentials: this.options.credentials,
+            credentials: this.params.credentials,
           })
         : new winston.transports.Console()
     );
@@ -74,7 +74,7 @@ export class GcpLogger implements LoggerConfig {
   }
 
   private shouldExport(env: string) {
-    return this.options.telemetryConfig?.forceDevExport || env !== 'dev';
+    return this.params.telemetryConfig?.forceDevExport || env !== 'dev';
   }
 }
 

js/plugins/google-cloud/src/gcpOpenTelemetry.ts

@@ -42,7 +42,7 @@ import {
   ReadableSpan,
   SpanExporter,
 } from '@opentelemetry/sdk-trace-base';
-import { PluginOptions } from './index.js';
+import { GoogleCloudPluginParams } from './index.js';
 
 let metricExporter: PushMetricExporter;
 let spanProcessor: BatchSpanProcessor;
@@ -53,7 +53,7 @@ let spanExporter: AdjustingTraceExporter;
  * Metrics, and Logs) to the Google Cloud Operations Suite.
  */
 export class GcpOpenTelemetry implements TelemetryConfig {
-  private readonly options: PluginOptions;
+  private readonly params: GoogleCloudPluginParams;
   private readonly resource: Resource;
 
   /**
@@ -63,14 +63,14 @@ export class GcpOpenTelemetry implements TelemetryConfig {
   private gcpTraceLogHook = (span: Span, record: any) => {
     const isSampled = !!(span.spanContext().traceFlags & TraceFlags.SAMPLED);
     record['logging.googleapis.com/trace'] = `projects/${
-      this.options.projectId
+      this.params.projectId
     }/traces/${span.spanContext().traceId}`;
     record['logging.googleapis.com/spanId'] = span.spanContext().spanId;
     record['logging.googleapis.com/trace_sampled'] = isSampled ? '1' : '0';
   };
 
-  constructor(options?: PluginOptions) {
-    this.options = options || {};
+  constructor(params?: GoogleCloudPluginParams) {
+    this.params = params || {};
     this.resource = new Resource({ type: 'global' }).merge(
       new GcpDetectorSync().detect()
     );
@@ -81,7 +81,7 @@ export class GcpOpenTelemetry implements TelemetryConfig {
     return {
       resource: this.resource,
       spanProcessor: spanProcessor,
-      sampler: this.options?.telemetryConfig?.sampler || new AlwaysOnSampler(),
+      sampler: this.params?.telemetryConfig?.sampler || new AlwaysOnSampler(),
       instrumentations: this.getInstrumentations(),
       metricReader: this.createMetricReader(),
     };
@@ -91,7 +91,7 @@ export class GcpOpenTelemetry implements TelemetryConfig {
     spanExporter = new AdjustingTraceExporter(
       this.shouldExportTraces()
         ? new TraceExporter({
-            credentials: this.options.credentials,
+            credentials: this.params.credentials,
           })
         : new InMemorySpanExporter()
     );
@@ -105,36 +105,36 @@ export class GcpOpenTelemetry implements TelemetryConfig {
     metricExporter = this.buildMetricExporter();
     return new PeriodicExportingMetricReader({
       exportIntervalMillis:
-        this.options?.telemetryConfig?.metricExportIntervalMillis || 300_000,
+        this.params?.telemetryConfig?.metricExportIntervalMillis || 300_000,
       exportTimeoutMillis:
-        this.options?.telemetryConfig?.metricExportTimeoutMillis || 300_000,
+        this.params?.telemetryConfig?.metricExportTimeoutMillis || 300_000,
       exporter: metricExporter,
     });
   }
 
   /** Gets all open telemetry instrumentations as configured by the plugin. */
   private getInstrumentations() {
-    if (this.options?.telemetryConfig?.autoInstrumentation) {
+    if (this.params?.telemetryConfig?.autoInstrumentation) {
       return getNodeAutoInstrumentations(
-        this.options?.telemetryConfig?.autoInstrumentationConfig || {}
+        this.params?.telemetryConfig?.autoInstrumentationConfig || {}
       ).concat(this.getDefaultLoggingInstrumentations());
     }
     return this.getDefaultLoggingInstrumentations();
   }
 
   private shouldExportTraces(): boolean {
     return (
-      (this.options.telemetryConfig?.forceDevExport ||
+      (this.params.telemetryConfig?.forceDevExport ||
         process.env.GENKIT_ENV !== 'dev') &&
-      !this.options.telemetryConfig?.disableTraces
+      !this.params.telemetryConfig?.disableTraces
     );
   }
 
   private shouldExportMetrics(): boolean {
     return (
-      (this.options.telemetryConfig?.forceDevExport ||
+      (this.params.telemetryConfig?.forceDevExport ||
         process.env.GENKIT_ENV !== 'dev') &&
-      !this.options.telemetryConfig?.disableMetrics
+      !this.params.telemetryConfig?.disableMetrics
     );
   }
 
@@ -149,12 +149,12 @@ export class GcpOpenTelemetry implements TelemetryConfig {
   private buildMetricExporter(): PushMetricExporter {
     const exporter: PushMetricExporter = this.shouldExportMetrics()
       ? new MetricExporter({
-          projectId: this.options.projectId,
+          projectId: this.params.projectId,
           userAgent: {
             product: 'genkit',
             version: GENKIT_VERSION,
           },
-          credentials: this.options.credentials,
+          credentials: this.params.credentials,
         })
       : new InMemoryMetricExporter(AggregationTemporality.DELTA);
     exporter.selectAggregation = (instrumentType: InstrumentType) => {