fix: get rid of some redundancy and fix problem with validating exist…

…ing files when using gemini
firebase · Mar 4, 2025 · 39efe53 · 39efe53
1 parent 292150e
commit 39efe53
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 31 deletions.
diff --git a/firestore-bigquery-export/scripts/gen-schema-view/src/config/index.ts b/firestore-bigquery-export/scripts/gen-schema-view/src/config/index.ts
@@ -12,8 +12,6 @@ export interface CliConfig {
   bigQueryProjectId: string;
   datasetId: string;
   tableNamePrefix: string;
-  // TODO: isn't this the same as tableNamePrefix? check.
-  collectionPath?: string;
   schemas: { [schemaName: string]: FirestoreSchema };
   useGemini?: boolean;
   agentSampleSize?: number;
@@ -33,9 +31,8 @@ export async function parseConfig(): Promise<CliConfig> {
       bigQueryProjectId: program.bigQueryProject || program.project,
       datasetId: program.dataset,
       tableNamePrefix: program.tableNamePrefix,
-      collectionPath: program.collectionPath,
-      schemas: readSchemas(program.schemaFiles),
       useGemini: program.useGemini,
+      schemas: !program.useGemini ? readSchemas(program.schemaFiles) : {},
       agentSampleSize: DEFAULT_SAMPLE_SIZE,
       googleAiKey: program.googleAiKey,
     };
@@ -46,7 +43,6 @@ export async function parseConfig(): Promise<CliConfig> {
     dataset,
     tableNamePrefix,
     schemaFiles,
-    collectionPath,
     useGemini,
     // TODO: rename?
     googleAiKey,
@@ -57,10 +53,9 @@ export async function parseConfig(): Promise<CliConfig> {
     bigQueryProjectId: bigQueryProject,
     datasetId: dataset,
     tableNamePrefix: tableNamePrefix,
-    collectionPath: collectionPath,
-    schemas: readSchemas(
+    schemas: !useGemini ? readSchemas(
       schemaFiles.split(",").map((schemaFileName) => schemaFileName.trim())
-    ),
+    ) : {},
     useGemini: useGemini,
     agentSampleSize: DEFAULT_SAMPLE_SIZE,
     googleAiKey: googleAiKey,

diff --git a/firestore-bigquery-export/scripts/gen-schema-view/src/config/interactive.ts b/firestore-bigquery-export/scripts/gen-schema-view/src/config/interactive.ts
@@ -48,29 +48,30 @@ export const questions = [
     validate: (value) =>
       validateInput(value, "table name prefix", BIGQUERY_VALID_CHARACTERS),
   },
-  {
-    message:
-      "Where should this script look for schema definitions? (Enter a comma-separated list of, optionally globbed, paths to files or directories).",
-    name: "schemaFiles",
-    type: "input",
-  },
   {
     message:
       "Would you like to use a Gemini to automatically analyze your data and generate a draft schema?",
     name: "useGemini",
     type: "confirm",
     default: false,
   },
-  // TODO: I dont think this is required as we have it above
-  // TODO: can we make the questions conditional? if we select useGemini then dont ask about finding schema files?
   {
-    message: "What is the Firestore collection path you want to analyze?",
-    name: "collectionPath",
+    message:
+      "Where should this script look for schema definitions? (Enter a comma-separated list of, optionally globbed, paths to files or directories).",
+    name: "schemaFiles",
     type: "input",
-    when: (answers) => answers.useGemini,
-    validate: (value) =>
-      validateInput(value, "collection path", FIRESTORE_VALID_CHARACTERS),
+    when: (answers) => !answers.useGemini,
   },
+  // TODO: I dont think this is required as we have it above
+  // TODO: can we make the questions conditional? if we select useGemini then dont ask about finding schema files?
+  // {
+  //   message: "What is the Firestore collection path you want to analyze?",
+  //   name: "collectionPath",
+  //   type: "input",
+  //   when: (answers) => answers.useGemini,
+  //   validate: (value) =>
+  //     validateInput(value, "collection path", FIRESTORE_VALID_CHARACTERS),
+  // },
   {
     message: "Please provide your Google AI API Key:",
     name: "googleAiKey",

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
@@ -194,12 +194,12 @@ const biqquerySchemaPrompt = ({
 export const generateSchemaFilesWithGemini = async (config: CliConfig) => {
   //  get sample data from Firestore
   const sampleData = await sampleFirestoreDocuments(
-    config.collectionPath!,
+    config.tableNamePrefix!,
     config.agentSampleSize!
   );
 
   const prompt = biqquerySchemaPrompt({
-    collectionName: config.collectionPath!,
+    collectionName: config.tableNamePrefix!,
     sampleData,
     tablePrefix: config.tableNamePrefix,
   });
@@ -214,21 +214,15 @@ export const generateSchemaFilesWithGemini = async (config: CliConfig) => {
   });
 
   // prompt gemini with sample data to generate a schema file
-  const { text, output } = await ai.generate({
+  const { text } = await ai.generate({
     model: gemini20Flash,
     prompt,
     output: {
       format: 'json',
-      schema: z.any()
+      schema: SchemaSchema
     }
   });
 
-  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([