[FDC] Generate schema during init dataconnect flow (#8596)

Author: fredzqm

CHANGELOG.md

@@ -0,0 +1,2 @@
+- Add `generate_dataconnect_schema`, `dataconnect_generate_operation`, `firebase_consult_assistant` MCP tools. (#8647)
+- `firebase init dataconnect` is now integrated with Gemini in Firebase API to generate Schema based on description. (#8596)

src/api.ts

@@ -91,8 +91,6 @@ export const functionsDefaultRegion = () =>
 
 export const cloudbuildOrigin = () =>
   utils.envOverride("FIREBASE_CLOUDBUILD_URL", "https://cloudbuild.googleapis.com");
-export const cloudCompanionOrigin = () =>
-  utils.envOverride("CLOUD_COMPANION_URL", "https://cloudaicompanion.googleapis.com");
 export const cloudschedulerOrigin = () =>
   utils.envOverride("FIREBASE_CLOUDSCHEDULER_URL", "https://cloudscheduler.googleapis.com");
 export const cloudTasksOrigin = () =>

src/dataconnect/ensureApis.ts

@@ -1,15 +1,19 @@
 import * as api from "../api";
 import { ensure } from "../ensureApiEnabled";
 
+const prefix = "dataconnect";
+
 export async function ensureApis(projectId: string): Promise<void> {
-  const prefix = "dataconnect";
   await ensure(projectId, api.dataconnectOrigin(), prefix);
   await ensure(projectId, api.cloudSQLAdminOrigin(), prefix);
   await ensure(projectId, api.computeOrigin(), prefix);
 }
 
 export async function ensureSparkApis(projectId: string): Promise<void> {
-  const prefix = "dataconnect";
   // These are the APIs that can be enabled without a billing account.
   await ensure(projectId, api.cloudSQLAdminOrigin(), prefix);
 }
+
+export async function ensureGIFApis(projectId: string): Promise<void> {
+  await ensure(projectId, api.cloudAiCompanionOrigin(), prefix);
+}

src/gemini/fdcExperience.spec.ts

@@ -0,0 +1,65 @@
+import { expect } from "chai";
+import { extractCodeBlock } from "./fdcExperience";
+
+describe("extractCodeBlock", () => {
+  it("should extract a basic GraphQL query block", () => {
+    const text =
+      'Here is a GraphQL query:\n```graphql\nquery GetUser { user(id: "1") { name email } }\n```\nThanks!';
+    const expected = 'query GetUser { user(id: "1") { name email } }';
+    expect(extractCodeBlock(text)).to.eq(expected);
+  });
+
+  it("should extract a multi-line GraphQL mutation block", () => {
+    const text = `
+    Some preamble.
+    \`\`\`graphql
+    mutation CreatePost($title: String!, $content: String!) {
+      createPost(title: $title, content: $content) {
+        id
+        title
+      }
+    }
+    \`\`\`
+    Followed by some description.
+    `;
+    const expected = `mutation CreatePost($title: String!, $content: String!) {
+      createPost(title: $title, content: $content) {
+        id
+        title
+      }
+    }`;
+    expect(extractCodeBlock(text)).to.eq(expected);
+  });
+
+  it("should extract a GraphQL fragment block", () => {
+    const text = "```graphql\nfragment UserFields on User { id name }\n```";
+    const expected = "fragment UserFields on User { id name }";
+    expect(extractCodeBlock(text)).to.eq(expected);
+  });
+
+  it("should extract an empty GraphQL code block", () => {
+    const text = "```graphql\n\n```";
+    const expected = "";
+    expect(extractCodeBlock(text)).to.eq(expected);
+  });
+
+  it("should extract a GraphQL schema definition block", () => {
+    const text = `
+    \`\`\`graphql
+    type Query {
+      hello: String
+    }
+    schema {
+      query: Query
+    }
+    \`\`\`
+    `;
+    const expected = `type Query {
+      hello: String
+    }
+    schema {
+      query: Query
+    }`;
+    expect(extractCodeBlock(text)).to.eq(expected);
+  });
+});

src/gemini/fdcExperience.ts

@@ -1,14 +1,15 @@
 import { Client } from "../apiv2";
-import { cloudCompanionOrigin } from "../api";
+import { cloudAiCompanionOrigin } from "../api";
 import {
   ChatExperienceResponse,
   CloudAICompanionMessage,
   CloudAICompanionRequest,
   GenerateOperationResponse,
   GenerateSchemaResponse,
 } from "./types";
+import { FirebaseError } from "../error";
 
-const apiClient = new Client({ urlPrefix: cloudCompanionOrigin(), auth: true });
+const apiClient = new Client({ urlPrefix: cloudAiCompanionOrigin(), auth: true });
 const SCHEMA_GENERATOR_EXPERIENCE = "/appeco/firebase/fdc-schema-generator";
 const GEMINI_IN_FIREBASE_EXPERIENCE = "/appeco/firebase/firebase-chat/free";
 const OPERATION_GENERATION_EXPERIENCE = "/appeco/firebase/fdc-query-generator";
@@ -91,3 +92,17 @@ export async function generateOperation(
   );
   return res.body.output.messages[0].content;
 }
+
+/**
+ * extractCodeBlock extracts the code block from the generated response.
+ * @param text the generated response from the service.
+ * @return the code block from the generated response.
+ */
+export function extractCodeBlock(text: string): string {
+  const regex = /```(?:[a-z]+\n)?([\s\S]*?)```/m;
+  const match = text.match(regex);
+  if (match && match[1]) {
+    return match[1].trim();
+  }
+  throw new FirebaseError(`No code block found in the generated response: ${text}`);
+}

src/init/features/dataconnect/index.ts

@@ -8,7 +8,7 @@ import { Setup } from "../..";
 import { provisionCloudSql } from "../../../dataconnect/provisionCloudSql";
 import { checkFreeTrialInstanceUsed, upgradeInstructions } from "../../../dataconnect/freeTrial";
 import * as cloudsql from "../../../gcp/cloudsql/cloudsqladmin";
-import { ensureApis, ensureSparkApis } from "../../../dataconnect/ensureApis";
+import { ensureApis, ensureGIFApis, ensureSparkApis } from "../../../dataconnect/ensureApis";
 import {
   listLocations,
   listAllServices,
@@ -19,10 +19,12 @@ import { Schema, Service, File, Platform } from "../../../dataconnect/types";
 import { parseCloudSQLInstanceName, parseServiceName } from "../../../dataconnect/names";
 import { logger } from "../../../logger";
 import { readTemplateSync } from "../../../templates";
-import { logBullet, logWarning, envOverride } from "../../../utils";
+import { logBullet, logWarning, envOverride, promiseWithSpinner } from "../../../utils";
 import { isBillingEnabled } from "../../../gcp/cloudbilling";
 import * as sdk from "./sdk";
 import { getPlatformFromFolder } from "../../../dataconnect/fileUtils";
+import { extractCodeBlock, generateSchema } from "../../../gemini/fdcExperience";
+import { configstore } from "../../../configstore";
 
 const DATACONNECT_YAML_TEMPLATE = readTemplateSync("init/dataconnect/dataconnect.yaml");
 const CONNECTOR_YAML_TEMPLATE = readTemplateSync("init/dataconnect/connector.yaml");
@@ -103,8 +105,7 @@ export async function askQuestions(setup: Setup): Promise<void> {
       default: true,
     }));
   if (shouldConfigureBackend) {
-    // TODO: Prompt for app idea and use GiF backend to generate them.
-    info = await promptForService(info);
+    info = await promptForSchema(setup, info);
     info = await promptForCloudSQL(setup, info);
 
     info.shouldProvisionCSQL = !!(
@@ -444,12 +445,38 @@ async function promptForCloudSQL(setup: Setup, info: RequiredInfo): Promise<Requ
   return info;
 }
 
-async function promptForService(info: RequiredInfo): Promise<RequiredInfo> {
+async function promptForSchema(setup: Setup, info: RequiredInfo): Promise<RequiredInfo> {
   if (info.serviceId === "") {
     info.serviceId = await input({
       message: "What ID would you like to use for this service?",
       default: basename(process.cwd()),
     });
+    if (setup.projectId) {
+      if (!configstore.get("gemini")) {
+        logBullet(
+          "Learn more about Gemini in Firebase and how it uses your data: https://firebase.google.com/docs/gemini-in-firebase#how-gemini-in-firebase-uses-your-data",
+        );
+      }
+      if (
+        await confirm({
+          message: `Do you want Gemini in Firebase to help generate a schema for your service?`,
+          default: false,
+        })
+      ) {
+        configstore.set("gemini", true);
+        await ensureGIFApis(setup.projectId);
+        const prompt = await input({
+          message: "Describe the app you are building:",
+          default: "movie rating app",
+        });
+        const schema = await promiseWithSpinner(
+          () => generateSchema(prompt, setup.projectId!),
+          "Generating the Data Connect Schema...",
+        );
+        info.schemaGql = [{ path: "schema.gql", content: extractCodeBlock(schema) }];
+        info.connectors = [emptyConnector];
+      }
+    }
   }
   return info;
 }