[VS Code] Add a code lens to refine / generate FDC operation (#9204)

Author: fredzqm

firebase-vscode/CHANGELOG.md

@@ -1,5 +1,7 @@
 ## NEXT
 
+- [Added] Refine / Generate Operation Code Lens.
+
 ## 1.8.0
 
 - [Changed] Gemini Code Assist is now optionally installed when using the "Build with AI" feature

firebase-vscode/src/analytics.ts

@@ -21,6 +21,7 @@ export enum DATA_CONNECT_EVENT_NAME {
   PROJECT_SELECTED = "project_selected",
   RUN_LOCAL = "run_local",
   RUN_PROD = "run_prod",
+  GENERATE_OPERATION = "generate_operation",
   ADD_DATA = "add_data",
   READ_DATA = "read_data",
   MOVE_TO_CONNECTOR = "move_to_connector",

firebase-vscode/src/data-connect/code-lens-provider.ts

@@ -6,6 +6,8 @@ import { Disposable } from "vscode";
 import { Signal } from "@preact/signals-core";
 import { dataConnectConfigs, firebaseRC } from "./config";
 import { EmulatorsController } from "../core/emulators";
+import { GenerateOperationInput } from "./execution/execution";
+import { findCommentsBlocks } from "../utils/find_comments";
 
 export enum InstanceType {
   LOCAL = "local",
@@ -85,6 +87,7 @@ export class OperationCodeLensProvider extends ComputedCodeLensProvider {
     for (let i = 0; i < documentNode.definitions.length; i++) {
       const x = documentNode.definitions[i];
       if (x.kind === Kind.OPERATION_DEFINITION && x.loc) {
+        // startToken.line is 1-indexed, range is 0-indexed
         const line = x.loc.startToken.line - 1;
         const range = new vscode.Range(line, 0, line, 0);
         const position = new vscode.Position(line, 0);
@@ -97,17 +100,6 @@ export class OperationCodeLensProvider extends ComputedCodeLensProvider {
           document.fileName,
         );
         if (service) {
-          // For demo purposes only
-          // codeLenses.push(
-          //   new vscode.CodeLens(range, {
-          //     title: `$(play) Refine Operation`,
-          //     command: "firebase.dataConnect.refineOperation",
-          //     tooltip:
-          //       "Execute the operation (⌘+enter or Ctrl+Enter)",
-          //     arguments: [x, operationLocation, InstanceType.LOCAL],
-          //   }),
-          // );
-
           codeLenses.push(
             new vscode.CodeLens(range, {
               title: `$(play) Run (local)`,
@@ -131,6 +123,33 @@ export class OperationCodeLensProvider extends ComputedCodeLensProvider {
       }
     }
 
+    if (projectId) {
+      const comments = findCommentsBlocks(documentText);
+      for (let i = 0; i < comments.length; i++) {
+        const c = comments[i];
+        const range = new vscode.Range(c.startLine, 0, c.startLine, 0);
+        const queryDoc = documentNode.definitions.find((d) =>
+          d.kind === Kind.OPERATION_DEFINITION && 
+          // startToken.line is 1-indexed, endLine is 0-indexed
+          d.loc?.startToken.line === c.endLine + 2
+        );
+        const arg: GenerateOperationInput = {
+          projectId,
+          document: document,
+          description: c.text,
+          insertPosition: c.endIndex + 1,
+          existingQuery: queryDoc?.loc ? documentText.substring(c.endIndex + 1, queryDoc.loc.endToken.end) : '',
+        };
+        codeLenses.push(
+          new vscode.CodeLens(range, {
+            title: queryDoc ? `$(sparkle) Refine Operation` : `$(sparkle) Generate Operation`,
+            command: "firebase.dataConnect.generateOperation",
+            tooltip: "Generate the operation (⌘+enter or Ctrl+Enter)",
+            arguments: [arg],
+          }),
+        );
+      }
+    }
     return codeLenses;
   }
 }

firebase-vscode/src/data-connect/execution/execution.ts

@@ -35,8 +35,11 @@ import { InstanceType } from "../code-lens-provider";
 import { DATA_CONNECT_EVENT_NAME, AnalyticsLogger } from "../../analytics";
 import { getDefaultScalarValue } from "../ad-hoc-mutations";
 import { EmulatorsController } from "../../core/emulators";
-import { getConnectorGQLText } from "../file-utils";
+import { getConnectorGQLText, insertQueryAt } from "../file-utils";
 import { pluginLogger } from "../../logger-wrapper";
+import * as gif from "../../../../src/gemini/fdcExperience";
+import { ensureGIFApiTos } from "../../../../src/dataconnect/ensureApis";
+import { configstore } from "../../../../src/configstore";
 
 interface TypedInput {
   varName: string;
@@ -49,6 +52,14 @@ interface ExecutionInput {
   instance: InstanceType;
 }
 
+export interface GenerateOperationInput {
+  projectId: string;
+  document: vscode.TextDocument;
+  description: string;
+  insertPosition: number;
+  existingQuery: string;
+}
+
 export const lastExecutionInputSignal = new Signal<ExecutionInput | null>(null);
 
 export function registerExecution(
@@ -302,6 +313,54 @@ export function registerExecution(
     }
   }
 
+  async function generateOperation(arg: GenerateOperationInput) {
+    try {
+      const schema = await dataConnectService.schema();
+      const prompt = `Generate a Data Connect operation to match this description: ${arg.description} 
+${arg.existingQuery ? `\n\nRefine this existing operation:\n${arg.existingQuery}` : ''}
+${schema ? `\n\nUse the Data Connect Schema:\n\`\`\`graphql
+${schema}
+\`\`\`` : ""}`;
+      const serviceName = await dataConnectService.servicePath(arg.document.fileName);
+      if (!(await ensureGIFApiTos(arg.projectId))) {
+        if (!(await showGiFToSModal(arg.projectId))) {
+          return; // ToS isn't accepted.
+        }
+      }
+      const res = await gif.generateOperation(prompt, serviceName, arg.projectId);
+      await insertQueryAt(arg.document.uri, arg.insertPosition, arg.existingQuery, res);
+    } catch (e: any) {
+      vscode.window.showErrorMessage(`Failed to generate query: ${e.message}`);
+    }
+  }
+
+  async function showGiFToSModal(projectId: string): Promise<boolean> {
+    const tos = "Terms of Service";
+    const enable = "Enable";
+    const result = await vscode.window.showWarningMessage(
+      "Gemini in Firebase",
+      {
+        modal: !process.env.VSCODE_TEST_MODE,
+        detail: "Gemini in Firebase helps you write Data Connect queries.",
+      },
+      enable,
+      tos,
+    );
+    switch (result) {
+      case enable:
+        configstore.set("gemini", true);
+        await ensureGIFApiTos(projectId);
+        return true;
+      case tos:
+        vscode.env.openExternal(
+          vscode.Uri.parse(
+            "https://firebase.google.com/docs/gemini-in-firebase#how-gemini-in-firebase-uses-your-data",
+          ),
+        );
+    }
+    return false;
+  }
+
   const sub4 = broker.on(
     "definedDataConnectArgs",
     (value) => (executionArgsJSON.value = value),
@@ -333,7 +392,16 @@ export function registerExecution(
             : DATA_CONNECT_EVENT_NAME.RUN_PROD,
         );
         await vscode.window.activeTextEditor?.document.save();
-        executeOperation(ast, location, instanceType);
+        await executeOperation(ast, location, instanceType);
+      },
+    ),
+    vscode.commands.registerCommand(
+      "firebase.dataConnect.generateOperation",
+      async (arg: GenerateOperationInput) => {
+        analyticsLogger.logger.logUsage(
+            DATA_CONNECT_EVENT_NAME.GENERATE_OPERATION,
+        );
+        await generateOperation(arg);
       },
     ),
     vscode.commands.registerCommand(

firebase-vscode/src/data-connect/file-utils.ts

@@ -4,7 +4,6 @@ import * as fs from "fs";
 
 import { dataConnectConfigs } from "./config";
 import { pluginLogger } from "../logger-wrapper";
-import { parse } from "graphql";
 
 export async function checkIfFileExists(file: Uri) {
   try {
@@ -56,25 +55,23 @@ export function getHighlightedText(): string {
   return editor.document.getText(selectionRange);
 }
 
-export function parseGraphql(content: string) {
-  content = content.replaceAll("```", "");
-  content = content.replaceAll("graphql", "");
-  const documentNode = parse(content);
-  return documentNode.definitions[0];
-}
-
-export function insertToBottomOfActiveFile(text: string) {
-  const editor = vscode.window.activeTextEditor;
-  if (!editor) {
+export async function insertQueryAt(uri: vscode.Uri, at: number, existing: string, replace: string): Promise<void> {
+  const doc = await vscode.workspace.openTextDocument(uri);
+  const text = doc.getText();
+  if (!existing) {
+    const newText = text.slice(0, at) + replace + text.slice(at);
+    await vscode.workspace.fs.writeFile(uri, new TextEncoder().encode(newText));
     return;
   }
-  const lastLine = editor.document.lineAt(editor.document.lineCount - 1);
-  const escapedText = text.replace(/\$/g, "\\\$"); // escape $ symbols
-
-  editor.insertSnippet(
-    new vscode.SnippetString(`\n\n${escapedText}`),
-    lastLine.range.end,
-  );
+  if (text.slice(at, at + existing.length) !== existing) {
+    throw new Error("The existing query was updated.");
+  }
+  // Adds a new line before if inserting at the end of the file
+  if (at > text.length) {
+    replace = "\n" + replace;
+  }
+  const newText = text.slice(0, at) + replace + text.slice(at + existing.length);
+  await vscode.workspace.fs.writeFile(uri, new TextEncoder().encode(newText));
 }
 
 // given a file path, compile all gql files for the associated connector

firebase-vscode/src/data-connect/service.ts

@@ -120,7 +120,7 @@ export class DataConnectService {
         operationName: "IntrospectionQuery",
         variables: "{}",
       });
-      console.log("introspection: ", introspectionResults);
+      console.log("introspection result: ", introspectionResults);
       // TODO: handle errors
       if ((introspectionResults as any).errors.length > 0) {
         return { data: undefined };
@@ -138,6 +138,23 @@ export class DataConnectService {
     }
   }
 
+  // Fetch the local Data Connect Schema sources via the toolkit introspection service.
+  async schema(): Promise<string> {
+    try {
+      const res = await this.executeGraphQLRead({
+        query: `query { _service { schema } }`,
+        operationName: "",
+        variables: "{}",
+      });
+      console.log("introspection schema result: ", res);
+      return (res as any)?.data?._service?.schema || "";
+    } catch (e) {
+      // TODO: surface error that emulator is not connected
+      pluginLogger.error("error: ", e);
+      return "";
+    }
+  }
+
   async executeGraphQLRead(params: {
     query: string;
     operationName: string;

firebase-vscode/src/utils/find_comments.ts

@@ -0,0 +1,37 @@
+interface Comment {
+  text: string;
+  startLine: number;
+  endLine: number;
+  endIndex: number;
+}
+
+export function findCommentsBlocks(text: string): Comment[] {
+  const lineEnds: number[] = [];
+  let searchIndex: number = -1;
+  while ((searchIndex = text.indexOf('\n', searchIndex + 1)) !== -1) {
+    lineEnds.push(searchIndex);
+  }
+  lineEnds.push(text.length);
+  const comments: Comment[] = [];
+  for (let i = 0; i < lineEnds.length; i++) {
+    const lineStart = i === 0 ? 0 : lineEnds[i - 1] + 1;
+    const lineText = text.substring(lineStart, lineEnds[i]).trim();
+    if (lineText.startsWith('#')) {
+      comments.push({ startLine: i, endLine: i, text: lineText.substring(1).trim(), endIndex: lineEnds[i] });
+    }
+  }
+  const commentBlocks: Comment[] = [];
+  for (let i = 0; i < comments.length; i++) {
+    const current = comments[i];
+    if (i === 0 || current.startLine > comments[i - 1].endLine + 1) {
+      commentBlocks.push({ ...current });
+    } else {
+      // Continuation of the previous block
+      const lastBlock = commentBlocks[commentBlocks.length - 1];
+      lastBlock.endLine = current.endLine;
+      lastBlock.endIndex = current.endIndex;
+      lastBlock.text += '\n' + current.text;
+    }
+  }
+  return commentBlocks;
+}

src/dataconnect/ensureApis.spec.ts

@@ -21,10 +21,4 @@ describe("ensureApis", () => {
     expect(ensureStub).to.be.calledWith("my-project", api.dataconnectOrigin(), "dataconnect");
     expect(ensureStub).to.be.calledWith("my-project", api.cloudSQLAdminOrigin(), "dataconnect");
   });
-
-  it("should ensure Cloud AI Companion API is enabled", async () => {
-    ensureStub.resolves();
-    await apis.ensureGIFApis("my-project");
-    expect(ensureStub).to.be.calledWith("my-project", api.cloudAiCompanionOrigin(), "dataconnect");
-  });
 });

src/dataconnect/ensureApis.ts

@@ -1,5 +1,6 @@
 import * as api from "../api";
-import { ensure } from "../ensureApiEnabled";
+import { configstore } from "../configstore";
+import { check, ensure } from "../ensureApiEnabled";
 
 const prefix = "dataconnect";
 
@@ -10,6 +11,18 @@ export async function ensureApis(projectId: string, silent: boolean = false): Pr
   ]);
 }
 
-export async function ensureGIFApis(projectId: string): Promise<void> {
-  await ensure(projectId, api.cloudAiCompanionOrigin(), prefix);
+/**
+ * Check if GIF APIs are enabled.
+ * If the Gemini in Firebase ToS is accepted, ensure the API is enabled.
+ * Otherwise, return false. The caller should prompt the user to accept the ToS.
+ */
+export async function ensureGIFApiTos(projectId: string): Promise<boolean> {
+  if (configstore.get("gemini")) {
+    await ensure(projectId, api.cloudAiCompanionOrigin(), "");
+  } else {
+    if (!(await check(projectId, api.cloudAiCompanionOrigin(), ""))) {
+      return false;
+    }
+  }
+  return true;
 }

src/init/features/dataconnect/index.ts

@@ -8,7 +8,7 @@ import { Setup } from "../..";
 import { setupCloudSql } from "../../../dataconnect/provisionCloudSql";
 import { checkFreeTrialInstanceUsed, upgradeInstructions } from "../../../dataconnect/freeTrial";
 import * as cloudsql from "../../../gcp/cloudsql/cloudsqladmin";
-import { ensureApis, ensureGIFApis } from "../../../dataconnect/ensureApis";
+import { ensureApis, ensureGIFApiTos } from "../../../dataconnect/ensureApis";
 import {
   listLocations,
   listAllServices,
@@ -118,7 +118,7 @@ export async function askQuestions(setup: Setup): Promise<void> {
       });
       if (info.appDescription) {
         configstore.set("gemini", true);
-        await ensureGIFApis(setup.projectId);
+        await ensureGIFApiTos(setup.projectId);
       }
     }
     if (hasBilling) {