fix: made some changes

firebase · Mar 4, 2025 · 292150e · 292150e
1 parent c7b9bb4
commit 292150e
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 31 deletions.
diff --git a/firestore-bigquery-export/scripts/gen-schema-view/src/index.ts b/firestore-bigquery-export/scripts/gen-schema-view/src/index.ts
@@ -39,20 +39,13 @@ async function run(): Promise<number> {
     config.bigQueryProjectId
   );
 
+  // Generate schema files using Gemini if enabled
+  // Otherwise, read schema files from the filesystem
+  let schemas = config.schemas;
   if (config.useGemini) {
     try {
       await generateSchemaFilesWithGemini(config);
-      const schemas = readSchemas([`./schemas/${config.tableNamePrefix}.json`]);
-
-      // TODO: move this out of the block so we're not repeating ourselves.
-      for (const name in schemas) {
-        await viewFactory.initializeSchemaViewResources(
-          config.datasetId,
-          config.tableNamePrefix,
-          name,
-          schemas[name]
-        );
-      }
+      schemas = readSchemas([`./schemas/${config.tableNamePrefix}.json`]);
 
       console.log("Schema views created successfully.");
     } catch (error) {
@@ -63,16 +56,18 @@ async function run(): Promise<number> {
     if (Object.keys(config.schemas).length === 0) {
       console.log(`No schema files found!`);
     }
+  }
 
-    for (const schemaName in config.schemas) {
-      await viewFactory.initializeSchemaViewResources(
-        config.datasetId,
-        config.tableNamePrefix,
-        schemaName,
-        config.schemas[schemaName]
-      );
-    }
+  // Initialize schema views
+  for (const name in schemas) {
+    await viewFactory.initializeSchemaViewResources(
+      config.datasetId,
+      config.tableNamePrefix,
+      name,
+      schemas[name]
+    );
   }
+
   return 0;
 }
 

diff --git a/firestore-bigquery-export/scripts/gen-schema-view/src/schema/genkit.ts b/firestore-bigquery-export/scripts/gen-schema-view/src/schema/genkit.ts
@@ -1,19 +1,11 @@
 import type { CliConfig } from "../config";
 import firebase = require("firebase-admin");
-import { genkit } from "genkit";
+import { genkit, z } from "genkit";
 import { googleAI, gemini20Flash } from "@genkit-ai/googleai";
 import * as fs from "fs/promises";
 import * as path from "path";
 import inquirer from "inquirer";
-
-const ai = genkit({
-  plugins: [
-    googleAI({
-      // TODO: we need to pass in the api key
-      // apiKey: config.googleAiKey,
-    }),
-  ],
-});
+import {SchemaSchema} from './genkitSchema'
 
 export async function sampleFirestoreDocuments(
   collectionPath: string,
@@ -194,7 +186,10 @@ const biqquerySchemaPrompt = ({
     ]
     }
   
-    Begin by analyzing the sample data and then create a well-documented schema.`;
+    Begin by analyzing the sample data and then create a well-documented schema.
+    
+    Please respond ONLY with the schema in json format
+    `;
 
 export const generateSchemaFilesWithGemini = async (config: CliConfig) => {
   //  get sample data from Firestore
@@ -209,12 +204,31 @@ export const generateSchemaFilesWithGemini = async (config: CliConfig) => {
     tablePrefix: config.tableNamePrefix,
   });
 
+  // initialize genkit with googleAI plugin
+  const ai = genkit({
+    plugins: [
+      googleAI({
+        apiKey: config.googleAiKey,
+      }),
+    ],
+  });
+
   // prompt gemini with sample data to generate a schema file
-  const { text } = await ai.generate({
+  const { text, output } = await ai.generate({
     model: gemini20Flash,
     prompt,
+    output: {
+      format: 'json',
+      schema: z.any()
+    }
   });
 
+  throw new Error(`gets to here ${JSON.stringify(output)}`)
+
+  console.log("this is output",output)
+
+  console.log(text);
+
   await writeSchemaFile("./schemas", `${config.tableNamePrefix}.json`, text);
   // confirm with user that schema file is correct
   const confirmation = await inquirer.prompt([