diff --git a/firestore-bigquery-export/scripts/gen-schema-view/src/index.ts b/firestore-bigquery-export/scripts/gen-schema-view/src/index.ts index e0520d407..359641483 100644 --- 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 { 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 { 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 index 9a55fd093..bd1f566d1 100644 --- 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([