diff --git a/firebase-vscode/src/data-connect/ai-tools/gca-tool-types.ts b/firebase-vscode/src/data-connect/ai-tools/gca-tool-types.ts deleted file mode 100644 index 183b81354f4..00000000000 --- a/firebase-vscode/src/data-connect/ai-tools/gca-tool-types.ts +++ /dev/null @@ -1,509 +0,0 @@ -import { - CancellationToken, - Disposable, - MarkdownString, - ThemeIcon, - Uri, -} from "vscode"; - -/** - * The public API for Gemini Code Assist to be utilized by external providers to - * extend Gemini Code Assist functionality. - */ -export interface GeminiCodeAssist extends Disposable { - /** - * Registers the caller as a tool for Gemini Code Assist. The tool will be - * identified to the end user through the id parameter. The tool will further - * identify itself through the extension id. An extension may choose to - * register any number of tools. - * @param id The id to use when referring to the tool. For example this may - * be `gemini` so that the tool will be addressed as `@gemini` by the user. - * Note that this id cannot be reused by another tool or other entity like - * a variable provider. - * @param displayName The name of the tool, to be used when referring to the - * tool in chat. - * @param extensionId The extension that implements the tool. The tool's - * icon will be loaded from this extension by default and used when - * displaying the tool's participation in chat. - * @param iconPath The path to the tool's icon, can be an icon for any theme, - * contain a dark and light icon or be a ThemeIcon type. The iconPath should be a join - * of the extension path and the relative path to the icon. - * @param command A command for Gemini Code Assist to execute on activation. - * If this is specified by the tool registration Gemini Code Assist will wait - * for the tool's extension to activate and then execute the command - * specified. This can be used to allow the tool to guarantee registration - * whenever Gemini Code Assist is loaded. - * @return The tool's registration to be modified by the tool provider with - * the capabilities of the tool. - */ - registerTool( - id: string, - displayName: string, - extensionId: string, - iconPath?: Uri | { dark: Uri; light: Uri } | ThemeIcon, - command?: string, - ): GeminiTool; -} - -/** - * Represents a tool to Gemini Code Assist. This allows the external provider - * to provide specific services to Gemini Code Assist. Upon dispose this tool - * registration will be removed from Gemini Code Assist for this instance only. - * Upon subsequent activations Gemini Code Assist will attempt to execute the - * command that was specified in the tool's registration if any was specified. - */ -export interface GeminiTool extends Disposable { - /** - * Registers a handler for chat. This allows the tool to handle incoming - * chat requests. - * @param handler The chat handler method that will be called with the - * registered tool is called. - * @return Disposable for subscription purposes, calling dispose will remove - * the registration. - */ - registerChatHandler(handler: ChatHandler): Disposable; - - /** - * Registers a variable provider for the tool. Variable provider ids should - * be unique for the tool but other tools may choose to implement the same - * provider id. For example `@bug` could be registered to `@jira` and - * `@github`. - * @param id The variable provider id, used to isolate typeahead to a specific - * variable type. For example using `@bug` will allow users to limit - * typeahead to bugs only instead of anything that can be completed. - * @param provider The provider to register, this will provide both static - * resolution as well as dynamic resolution. - * @return Disposable for removing the variable provider from Gemini Code - * Assist. - */ - registerVariableProvider(id: string, provider: VariableProvider): Disposable; - - /** - * Registers a slash command provider for the tool. This allows the tool - * to provide slash commands for the user. For example `/list` can be - * registered to the `@jira` tool to list bugs assigned to the user. - * @param provider The slash command provider to be registered. - * @return Disposable for removing the command provider from Gemini Code - * Assist. - */ - registerCommandProvider(provider: CommandProvider): Disposable; - - /** - * Registers a suggested prompt provider for the tool. This allows the tool - * to provide suggestions of prompts that the user can either tab complete or - * click to use. - * @param provider The child provider to be registered. - * @return Disposable for removing the suggested prompt provider from Gemini - * Code Assist. - */ - registerSuggestedPromptProvider( - provider: SuggestedPromptProvider, - ): Disposable; -} - -/** - * Provides suggested prompts which serve as example queries for the tool. - * These suggested prompts allow the user to see specific examples when using - * the tool and give some guidance as to helpful prompts as a starting point for - * using the tool with Gemini Code Assist. - */ -export interface SuggestedPromptProvider { - /** - * Provides a list of suggested prompts for the tool that will be displayed to - * the user as examples or templates for using the tool. In this text the - * user can specify placeholder text as text surrounded by square brackets. - * For example a suggested prompt value of `/generate an api specification for - * [function]` provided by `apigee` would provide a suggested prompt of - * `@apigee /generate an api specification for [function]` and the user would - * be prompted to supply a value for the [function] placeholder. - */ - provideSuggestedPrompts(): string[]; -} - -/** - * Provides the chat handler functionality to Gemini Code Assist, allowing a - * tool to extend chat. Through this handler the tool can service chat - * requests, add context for Gemini chat requests, and/or rewrite the prompt - * before it is sent to the LLM service. - */ -export interface ChatHandler { - /** - * @param request The chat request, can be used to manipulate the prompt and - * reference parts. - */ - ( - request: ChatRequest, - responseStream: ChatResponseStream, - token: CancellationToken, - ): Promise; -} - -/** - * Provides support for variables through the `@` designator. For example - * `@repo` could represent the current repository for a SCM tool. This - * interface allows the tool to provide a static list of variables as well as - * a dynamic list. - */ -export interface VariableProvider { - /** - * Allows the tool to return a static list of variables that it supports. - * This list is not expected to change as the user is typing. - * @return Returns a list of variables instances. - */ - listVariables(): Promise; - - /** - * Allows for dynamic variable support. This function will allow the tool - * to resolve variables as the user types. - * @param part Current text part that the user has typed, this is what - * currently follows the `@` symbol in the user's prompt. - * @param limit The number of typeahead suggestions that the UI will show to - * the user at once. - * @param token Supports cancellation (user types an additional character). - * @return Returns a list of variable instances that match the type ahead. - */ - typeahead( - part: string, - limit: number, - token: CancellationToken, - ): Promise; -} - -/** - * Represents a variable instance, the name and description are used to display - * the variable to the user. The variable instance will be passed as is to the - * tool, so it can carry any additional context necessary. - */ -export interface Variable { - /** - * The name of the variable, this would be what the variable looks like to the - * user. - */ - name: string; - - /** - * The optional description of the variable to show the user in the UX. - */ - description?: string | MarkdownString; -} - -/** - * Provides support for commands through the `/` designator. This takes the - * form of `@tool /command`. - */ -export interface CommandProvider { - /** - * Lists the slash commands provided by the tool. - * @return Command gives a list of the commands provided by the tool. - */ - listCommands(): Promise; -} - -/** - * CommandDetail exports a command along with any other context the tool may - * want to have in coordination with the command. - */ -export interface CommandDetail { - /** - * The string that identifies the command in question. - */ - command: string; - - /** - * The optional description of the slash command to display to the user. - */ - description?: string | MarkdownString; - - /** - * The optional codicon of the slash command. - */ - icon?: string; -} - -/** - * CommandPromptPart is the part of the prompt that is associated with a slash - * command. - */ -export interface CommandPromptPart extends PromptPart { - /** - * The CommandDetail provided by the CommandProvider's listCommands() - * function. - */ - command: CommandDetail; -} - -/** - * Provides the context for the chat request. The context can be used to - * provide additional information to the LLM service. - */ -export interface ChatRequestContext { - /** - * Pushes a new context onto the context stack. - * @param context The context to push. - */ - push(context: ChatContext | VariableChatContext): void; -} - -/** - * Represents a context that can be used to provide additional information to the - * LLM service. - */ -export interface ChatContext { - /** - * The id of the reference that this context is associated with. - */ - id: string | Uri; - - /** - * Gets the text of the context. - */ - getText(): string; -} - -/** - * Represents a context for a variable in the prompt. - */ -export interface VariableChatContext extends ChatContext { - /** - * The variable that this context represents. - */ - variable: Variable; -} - -/** - * Represents a chat request which is comprised of a prompt and context. - */ -export interface ChatRequest { - /** - * The prompt of the chat request. This can be manipulated by the tool. - */ - prompt: ChatPrompt; - - /** - * The context for the request. This can be used by the tool to add context - * to the request. - */ - context: ChatRequestContext; -} - -/** - * Represents the current chat prompt. - */ -export interface ChatPrompt { - /** - * Used to retrieve all parts of the prompt, including the tool prompts. - * @return An array of all parts of the prompt in order that they appear. - */ - getPromptParts(): PromptPart[]; - - /** - * Removes the specified prompt part. - * @param part The prompt part to remove. - */ - deletePromptPart(part: PromptPart): void; - - /** - * Splices in prompt part(s) similarly to Array.splice(). This can be used to - * insert a number of prompt part(s) (including none) and can remove existing - * elements. - * @param index The starting index for the splice operation. - * @param remove The number of elements to remove. - * @param parts The prompt part(s) to insert. - */ - splice(index: number, remove: number, ...parts: PromptPart[]): void; - - /** - * Pushes the prompt part(s) into the chat prompt. These part(s) are appended - * similarly to array.push(). - * @param parts The prompt part(s) to push. - */ - push(...parts: PromptPart[]): void; - - /** - * Returns the string representation of the prompt. - */ - fullPrompt(): string; - - /** - * The length of the prompt in parts. - */ - length: number; -} - -/** - * Represents a prompt part that is provided by a tool. - */ -export interface PromptPart { - /** - * Gets the prompt of the prompt part. - */ - getPrompt(): string; -} - -/** - * Represents a prompt part that is provided by a tool. - */ -export interface ToolPromptPart extends PromptPart { - /** - * The id of the tool that provided the prompt part. - */ - toolId: string; - - /** - * The command of the prompt part. - */ - command: string; -} - -/** - * Represents a prompt part that refers to a variable. - */ -export interface VariablePromptPart extends PromptPart { - variable: Variable; -} - -/** - * Represents a stream of chat responses. Used by the tool to provide chat - * based responses to the user. This stream can be used to push both partial - * responses as well as to close the stream. - */ -export interface ChatResponseStream { - /** - * Pushes a new content onto the response stream. - * @param content The content to push. - */ - push(content: MarkdownString | Citation): void; - - /** - * Closes the steam and prevents the request from going to the LLM after tool - * processing. This can be utilized by the client for commands that are - * client only. Returning without calling close will result in the processed - * prompt and context being sent to the LLM for a result. - */ - close(): void; - - /** - * Adds a button handler to code responses that come back from the LLM. This - * allows the tool to present the user with a button attached to this response - * and on click process the code response. - * @param title The title of the button. - * @param handler The handler to execute when the user clicks on the button. - * The code block will be sent as an argument to the handler on execution as - * CodeHandlerCommandArgs. - * @param languageFilter Optional parameter, if this is specified the - * language specified on the block will be checked for a match against this - * and the button will be displayed on match. If this is not specified the - * buttone will be displayed on any language result. - */ - addCodeHandlerButton( - title: string, - handler: CodeHandler, - options: HandlerButtonOptions, - ): void; -} - -/** - * Method for handling code responses from the LLM attached to the response via - * ChatResponseStream.addCodeHandlerButton. - */ -export interface CodeHandler { - /** - * @param args The code block and language specifiers from the LLM response to - * handle. Called when the user clicks on a CodeHandlerButton. - */ - (args: CodeHandlerCommandArgs): void; -} - -/** - * The arguments that are sent when calling the command associated with the - * addCodeHandlerButton method. - */ -export interface CodeHandlerCommandArgs { - /** - * The code block that was attached to the code handler button. - */ - codeBlock: string; - /** - * The language specifier on the code block associated with the code handler - * button. - */ - language: string; -} - -/** - * Provides options for a code handler button. This allows the tool to - * specialize the way GCA handles specific code blocks when the agent is - * involved. - */ -export interface HandlerButtonOptions { - /** - * Optional parameter, if this is specified the language specified on the - * block will be checked for a match against this and the button will be - * displayed on match. If this is not specified the button will be displayed - * on any language result. - */ - languages?: RegExp; - - /** - * Optional, if this is specified the block will be either expanded or - * collapsed as specified. If this is not specified the built in default - * handler will be used. - */ - displayType?: BlockDisplayType; -} - -/** - * Specifies how code blocks should be handled in chat. - */ -export enum BlockDisplayType { - /** - * The code block will be expanded by default. - */ - Expanded, - - /** - * The code block will be collapsed by default. - */ - Collapsed, -} - -/** - * Represents the type of citation. - */ -export enum CitationType { - /** - * General citation where the link is from a unspecific source. - */ - Unknown, - - /** - * The citation originates from the user's machine, for example a file on the - * user's disk. - */ - Local, - - /** - * The citation comes from Github. - */ - Github, -} - -/** - * Represents a citation. - */ -export interface Citation { - /** - * The URI of the citation. - */ - uri: Uri; - - /** - * The license of the citation. - */ - license: string | undefined; - - /** - * The type of the citation. - */ - type: CitationType; -} diff --git a/firebase-vscode/src/data-connect/ai-tools/gca-tool.ts b/firebase-vscode/src/data-connect/ai-tools/gca-tool.ts deleted file mode 100644 index e93554a5802..00000000000 --- a/firebase-vscode/src/data-connect/ai-tools/gca-tool.ts +++ /dev/null @@ -1,238 +0,0 @@ -import { AnalyticsLogger } from "../../analytics"; -import { ExtensionBrokerImpl } from "../../extension-broker"; -import * as vscode from "vscode"; -import { DataConnectService } from "../service"; -import { - ChatPrompt, - ChatRequest, - ChatResponseStream, - CommandDetail, - CommandProvider, - GeminiCodeAssist, -} from "./gca-tool-types"; -import { insertToBottomOfActiveFile } from "../file-utils"; -import { ExtensionContext } from "vscode"; -import { Chat, Command } from "./types"; -import { GeminiToolController } from "./tool-controller"; -import { ChatMessage } from "../../dataconnect/cloudAICompanionTypes"; -export const DATACONNECT_TOOL_ID = "FirebaseDataConnect"; -const AT_DATACONNECT_TOOL_ID = `@${DATACONNECT_TOOL_ID}`; -export const DATACONNECT_DISPLAY_NAME = "Firebase Data Connect"; -export const SUGGESTED_PROMPTS = [ - "/generate_schema Create a schema for a pizza store", - "/generate_operation Create a mutations for all my types", -]; -const HELP_MESSAGE = ` -Welcome to the Data Connect Tool. -Usage: - ${AT_DATACONNECT_TOOL_ID} /generate_schema \n - ${AT_DATACONNECT_TOOL_ID} /generate_operation -`; - -export class GCAToolClient { - private history: Chat[] = []; - private icon = vscode.Uri.joinPath( - this.context.extensionUri, - "resources", - "firebase_dataconnect_logo.png", - ); - constructor( - private context: ExtensionContext, - private toolController: GeminiToolController, - ) {} - - async activate() { - const gemini = vscode.extensions.getExtension( - "google.geminicodeassist", - ); - if (!gemini || !gemini.isActive) { - throw new Error("Gemini extension not found"); // should never happen, gemini is an extension depedency - } - - gemini?.activate().then(async (gca) => { - const tool = gca.registerTool( - DATACONNECT_TOOL_ID, - DATACONNECT_DISPLAY_NAME, - "GoogleCloudTools.firebase-dataconnect-vscode", - this.icon, - "help", - ); - tool.registerChatHandler(this.handleChat.bind(this)); - tool.registerSuggestedPromptProvider(this); - tool.registerCommandProvider( - new DataConnectCommandProvider(this.icon.toString()), - ); - }); - } - - /** implementation of handleChat interface; - * We redirect the request to our controller - */ - async handleChat( - request: ChatRequest, - responseStream: ChatResponseStream, - token: vscode.CancellationToken, - ): Promise { - // Helper just to convert to markdown first - function pushToResponseStream(text: string) { - const markdown = new vscode.MarkdownString(text); - responseStream.push(markdown); - } - - // Adds the Graphql code block button "Insert to bottom of file" - addCodeHandlers(responseStream); - - let response: ChatMessage[]; - - // parse the prompt - if (!isPromptValid(request.prompt)) { - pushToResponseStream(HELP_MESSAGE); - responseStream.close(); - return; - } - const content = getPrompt(request.prompt); - const command = getCommand(request.prompt); - - // Forward to tool controller - try { - this.history.push({ author: "USER", content, commandContext: command }); - response = await this.toolController.handleChat( - content, - this.history, - command, - ); - } catch (error) { - let errorMessage = ""; - if (error instanceof Error) { - errorMessage = error.message; - } else if (typeof error === "string") { - errorMessage = error; - } - - pushToResponseStream(errorMessage); - - // reset history on error - this.history = []; - responseStream.close(); - return; - } - const agentMessage = response.pop()?.content; - - if (agentMessage) { - this.history.push({ author: "AGENT", content: agentMessage }); - } - - pushToResponseStream( - agentMessage || "Gemini encountered an error. Please try again.}", - ); - responseStream.close(); - } - - provideSuggestedPrompts(): string[] { - return SUGGESTED_PROMPTS; - } -} - -class DataConnectCommandProvider implements CommandProvider { - schemaCommand: CommandDetail = { - command: Command.GENERATE_SCHEMA, - description: "Generates a GraphQL schema based on a prompt", - icon: this.icon, - }; - - operationCommand: CommandDetail = { - command: Command.GENERATE_OPERATION, - description: "Generates a GraphQL query or mutation based on a prompt", - icon: this.icon, - }; - - helpCommand: CommandDetail = { - command: "help", - description: "Shows this help message", - icon: this.icon, - }; - constructor(readonly icon: string) {} - listCommands(): Promise { - const commands: CommandDetail[] = [ - this.schemaCommand, - this.operationCommand, - // this.helpCommand, - ]; - return Promise.resolve(commands); - } -} - -/** Exploring a variable provider for dataconnect introspected types */ -// class DataConnectTypeVariableProvider implements VariableProvider { -// constructor(private fdcService: DataConnectService) {} -// async listVariables(): Promise { -// const introspection = await this.fdcService.introspect(); -// console.log(introspection); -// return introspection.data!.__schema.types.map((type) => { -// return { -// name: type.name, -// description: type.description as string, -// }; -// }); -// } - -// typeahead( -// part: string, -// limit: number, -// token: vscode.CancellationToken, -// ): Promise { -// throw new Error("Method not implemented."); -// } -// } - -// currently only supports a single button -function addCodeHandlers(responseStream: ChatResponseStream) { - responseStream.addCodeHandlerButton( - "Insert to bottom of file", - ({ codeBlock }) => { - insertToBottomOfActiveFile(codeBlock); - }, - { languages: /graphql|graphqllanguage/ }, - ); -} - -// Basic validation function to ensure deterministic command -function isPromptValid(prompt: ChatPrompt): boolean { - if (prompt.length < 2) { - return false; - } - if (prompt.getPromptParts()[0].getPrompt() !== AT_DATACONNECT_TOOL_ID) { - return false; - } - - return isCommandValid( - prompt.getPromptParts()[1].getPrompt().replace("/", ""), - ); -} - -function isCommandValid(command: string): boolean { - return (Object.values(Command) as string[]).includes(command); -} - -// get the /command without the / -function getCommand(prompt: ChatPrompt): Command { - if (prompt.length > 2) { - return prompt.getPromptParts()[1].getPrompt().replace("/", "") as Command; - } - - // fallback if prompt parts doesn't work - return prompt.fullPrompt().replace(AT_DATACONNECT_TOOL_ID, "").trimStart().split(" ")[0] as Command; -} - -// get the entire prompt without the @tool & /command -function getPrompt(prompt: ChatPrompt): string { - if ( - prompt.length > 2 && - prompt.getPromptParts()[0].getPrompt() === AT_DATACONNECT_TOOL_ID - ) { - return prompt.getPromptParts()[2].getPrompt(); - } - - // fallback if prompt parts doesn't work - return prompt.fullPrompt().replace(AT_DATACONNECT_TOOL_ID, "").replace(/\/\w+/, "").trimStart(); -} diff --git a/firebase-vscode/src/data-connect/ai-tools/tool-controller.ts b/firebase-vscode/src/data-connect/ai-tools/tool-controller.ts deleted file mode 100644 index 500334099b6..00000000000 --- a/firebase-vscode/src/data-connect/ai-tools/tool-controller.ts +++ /dev/null @@ -1,289 +0,0 @@ -import * as fs from "fs"; -import * as path from "path"; -import * as vscode from "vscode"; -import { Signal } from "@preact/signals-core"; - -import { Result } from "../../result"; -import { AnalyticsLogger } from "../../analytics"; -import { ResolvedDataConnectConfigs } from "../config"; -import { DataConnectService } from "../service"; -import { CloudAICompanionResponse, ChatMessage } from "../../dataconnect/cloudAICompanionTypes"; -import { ObjectTypeDefinitionNode, OperationDefinitionNode } from "graphql"; -import { getHighlightedText, findGqlFiles } from "../file-utils"; -import { CommandContext, Chat, Context, Command, BackendAuthor } from "./types"; -import { DATA_CONNECT_EVENT_NAME } from "../../analytics"; - -const USER_PREAMBLE = "This is the user's prompt: \n"; - -const SCHEMA_PROMPT_PREAMBLE = - "This is the user's current schema in their code base.: \n"; - -const NEW_LINE = "\n"; -const HIGHLIGHTED_TEXT_PREAMBLE = - "This is the highlighted code in the users active editor: \n"; - -/** - * Logic for talking to CloudCompanion API - * Handles Context collection and management - * - */ -export class GeminiToolController { - constructor( - private readonly analyticsLogger: AnalyticsLogger, - private readonly fdcService: DataConnectService, - private configs: Signal< - Result | undefined - >, - ) { - this.registerCommands(); - } - - // entry points from vscode to respsective tools - private registerCommands(): void { - /** Demo only */ - // vscode.commands.registerCommand( - // "firebase.dataConnect.refineOperation", - // async (ast: ObjectTypeDefinitionNode) => { - // this.highlightActiveType(ast); - // if (env.value.isMonospace) { - // vscode.commands.executeCommand("aichat.prompt", { - // prefillPrompt: "@data-connect /generate_operation ", - // }); - // } else { - // // change to prefill when GCA releases feature - // vscode.commands.executeCommand("cloudcode.gemini.chatView.focus"); - // } - // }, - // ); - /** End Demo only */ - } - private highlightActiveType(ast: ObjectTypeDefinitionNode) { - const editor = vscode.window.activeTextEditor; - if (!editor || !ast.loc) { - // TODO: add a warning, and skip this process - } else { - // highlight the schema in question - const startPostion = new vscode.Position( - ast.loc?.startToken.line - 1, - ast.loc?.startToken.column - 1, - ); - const endPosition = new vscode.Position( - ast.loc?.endToken.line, - ast.loc?.endToken.column - 1, - ); - editor.selection = new vscode.Selection(startPostion, endPosition); - } - } - - /** - * Entry point to chat interface; - * Builds prompt given chatHistory and generation type - * We use some basic heuristics such as - * - presence of previously generated code - * - activeEditor + any highlighted code - */ - public async handleChat( - userPrompt: string, // prompt without toolname and command - chatHistory: Chat[], - command: Command, - ): Promise { - let prompt = ""; - let currentChat: Chat = { - author: "USER", - content: "to_be_set", - commandContext: CommandContext.NO_OP /* to be set */, - }; - let type: "schema" | "operation"; - - // set type - if (command === Command.GENERATE_OPERATION) { - type = "operation"; - this.analyticsLogger.logger.logUsage( - DATA_CONNECT_EVENT_NAME.GEMINI_OPERATION_CALL, - ); - } else if (command === Command.GENERATE_SCHEMA) { - type = "schema"; - this.analyticsLogger.logger.logUsage( - DATA_CONNECT_EVENT_NAME.GEMINI_SCHEMA_CALL, - ); - } else { - // undetermined process - chatHistory.push({ - author: "MODEL", - content: - "Gemini is unable to complete that request. Try '/generate_schema' or '/generate_operation' to get started.", - }); - return chatHistory; - } - - //TODO: deal with non-open editor situation - const currentDocumentPath = - vscode.window.activeTextEditor?.document.uri.path; - - // get additional context - const schema = await this.collectSchemaText(); - const highlighted = getHighlightedText(); - - // check if highlighted is a single operation - if (highlighted) { - prompt = prompt.concat(HIGHLIGHTED_TEXT_PREAMBLE, highlighted); - } - - // only add schema for operation generation - if (schema && command === Command.GENERATE_OPERATION) { - prompt = prompt.concat(SCHEMA_PROMPT_PREAMBLE, schema); - } - - // finalize prompt w/ user prompt - prompt = prompt.concat(USER_PREAMBLE, userPrompt); - - const resp = await this.callGenerateApi( - currentDocumentPath || "", - prompt, - type, - this.cleanHistory(chatHistory, type), - ); - - if (resp.error) { - this.analyticsLogger.logger.logUsage( - DATA_CONNECT_EVENT_NAME.GEMINI_ERROR, - ); - return [{ author: "MODEL", content: resp.error.message }]; - } - - return resp.output.messages; - } - - // clean history for API consumption - public cleanHistory(history: ChatMessage[], type: string): Chat[] { - if (type === "operation") { - // operation api uses "SYSTEM" to represent API responses - return history.map((item) => { - if ( - item.author.toUpperCase() === "MODEL" || - item.author.toUpperCase() === "AGENT" - ) { - item.author = "SYSTEM"; - } - - if (item.author.toUpperCase() === "USER") { - item.author = "USER"; // set upper case - } - // remove command context - return { author: item.author, content: item.content }; - }); - } else { - return history.map((item) => { - if ( - item.author.toUpperCase() === "AGENT" || - item.author.toUpperCase() === "SYSTEM" - ) { - item.author = "MODEL"; - } - item.author = item.author.toUpperCase(); - - return { - author: item.author, - content: item.content, - }; - }); - } - } - - async callGenerateApi( - documentPath: string, - prompt: string, - type: "schema" | "operation", - chatHistory: Chat[], - ): Promise { - // TODO: Call Gemini API with the document content and context - try { - const response = await this.fdcService.generateOperation( - documentPath, - prompt, - type, - chatHistory, - ); - if (!response) { - throw new Error("No response from Cloud AI API"); - } - return response; - } catch (error) { - throw new Error(`Failed to call Gemini API: ${error}`); - } - } - - async collectSchemaText(): Promise { - try { - const service = this.configs?.value?.tryReadValue?.values[0]; - - if (!service) { - // The entrypoint is not a codelens file, so we can't determine the service. - return ""; - } - - let schema: string = ""; - const schemaPath = path.join(service.path, service.schemaDir); - const schemaFiles = await findGqlFiles(schemaPath); - for (const file of schemaFiles) { - schema = schema.concat(fs.readFileSync(file, "utf-8")); - } - return schema; - } catch (error) { - throw new Error(`Failed to collect GQL files: ${error}`); - } - } - - /** Demo usage only */ - private async setupRefineOperation(prompt: string, chatHistory: Chat[]) { - const preamble = - "This is the GraphQL Operation that was generated previously: "; - - // TODO: more verification - const lastChat = chatHistory.pop(); - let operation = ""; - if (!lastChat) { - // could not find an operation, TODO: response appropriately - } else { - operation = lastChat.content; - } - - return preamble.concat(NEW_LINE, operation); - } - - private async setupRefineSchema(prompt: string, chatHistory: Chat[]) { - const SCHEMA_PREAMBLE = - "This is the GraphQL Schema that was generated previously: \n"; - - // TODO: more verification - const lastChat = chatHistory.pop(); - let schema = ""; - if (!lastChat) { - // could not find a schema, use the schema in editor - schema = await this.collectSchemaText(); - } else { - schema = lastChat.content; - } - - return prompt.concat(SCHEMA_PREAMBLE, schema); - } - - private isAuthorBackend(author: string) { - return Object.values(BackendAuthor).includes(author); - } - - // checks if last chat in the history is a generated code response from a model - private isLastChatGenerated(chatHistory: Chat[]): boolean { - const lastChat = chatHistory.pop(); - return ( - lastChat !== undefined && - this.isAuthorBackend(lastChat.author) && - lastChat.commandContext !== undefined && - lastChat.commandContext !== CommandContext.NO_OP - ); - } - - /** End demo code */ - - dispose() {} -} diff --git a/firebase-vscode/src/data-connect/ai-tools/types.ts b/firebase-vscode/src/data-connect/ai-tools/types.ts deleted file mode 100644 index 8d3dbc7df29..00000000000 --- a/firebase-vscode/src/data-connect/ai-tools/types.ts +++ /dev/null @@ -1,25 +0,0 @@ -import { ChatMessage } from "../../dataconnect/cloudAICompanionTypes"; - -export enum Command { - GENERATE_SCHEMA = "generate_schema", - GENERATE_OPERATION = "generate_operation", -} -export enum Context { - REFINE_SCHEMA = "refine_schema", - REFINE_OPERATION = "refine_op", - NO_OP = "no_op", // not no_operation, it's just a no-op -} -// export type CommandContext = Command | Context; -export const CommandContext = { ...Command, ...Context }; -export type CommandContextType = Command | Context; - -// adds context to the ChatMessage type for reasoning -export interface Chat extends ChatMessage { - commandContext?: CommandContextType; -} - -// represents a backend chat response -export const BackendAuthor = { - MODEL: "MODEL", // schema api - SYSTEM: "SYSTEM", // operation api -}; diff --git a/firebase-vscode/src/data-connect/index.ts b/firebase-vscode/src/data-connect/index.ts index 8fc1bf1fcd1..98d2b7279b3 100644 --- a/firebase-vscode/src/data-connect/index.ts +++ b/firebase-vscode/src/data-connect/index.ts @@ -1,4 +1,4 @@ -import vscode, { Disposable, ExtensionContext, TelemetryLogger } from "vscode"; +import vscode, { Disposable, ExtensionContext } from "vscode"; import { Signal, effect } from "@preact/signals-core"; import { ExtensionBrokerImpl } from "../extension-broker"; import { registerExecution } from "./execution/execution"; @@ -31,14 +31,8 @@ import { registerWebview } from "../webview"; import { DataConnectToolkit } from "./toolkit"; import { registerFdcSdkGeneration } from "./sdk-generation"; import { registerDiagnostics } from "./diagnostics"; -import { AnalyticsLogger, DATA_CONNECT_EVENT_NAME } from "../analytics"; -import { emulators } from "../init/features"; -import { GCAToolClient } from "./ai-tools/gca-tool"; -import { GeminiToolController } from "./ai-tools/tool-controller"; -import { - registerFirebaseMCP, - writeToGeminiConfig, -} from "./ai-tools/firebase-mcp"; +import { AnalyticsLogger } from "../analytics"; +import { registerFirebaseMCP } from "./ai-tools/firebase-mcp"; class CodeActionsProvider implements vscode.CodeActionProvider { constructor( @@ -238,7 +232,6 @@ export function registerFdc( registerTerminalTasks(broker, analyticsLogger), registerFirebaseMCP(broker, analyticsLogger), operationCodeLensProvider, - vscode.languages.registerCodeLensProvider( // **Hack**: For testing purposes, enable code lenses on all graphql files // inside the test_projects folder. diff --git a/firebase-vscode/src/data-connect/service.ts b/firebase-vscode/src/data-connect/service.ts index 99eafe1d768..dd07679195f 100644 --- a/firebase-vscode/src/data-connect/service.ts +++ b/firebase-vscode/src/data-connect/service.ts @@ -10,7 +10,7 @@ import { AuthService } from "../auth/service"; import { UserMockKind } from "../../common/messaging/protocol"; import { firstWhereDefined } from "../utils/signal"; import { EmulatorsController } from "../core/emulators"; -import { dataConnectConfigs, VSCODE_ENV_VARS } from "../data-connect/config"; +import { dataConnectConfigs } from "../data-connect/config"; import { firebaseRC } from "../core/config"; import { @@ -20,22 +20,12 @@ import { DATACONNECT_API_VERSION, } from "../../../src/dataconnect/dataplaneClient"; -import { - cloudAICompationClient, - callCloudAICompanion, -} from "../../../src/dataconnect/cloudAiCompanionClient"; - import { ExecuteGraphqlRequest, GraphqlResponse, GraphqlResponseError, Impersonation, } from "../dataconnect/types"; -import { - CloudAICompanionResponse, - CallCloudAiCompanionRequest, - ChatMessage, -} from "../dataconnect/cloudAICompanionTypes"; import { Client, ClientResponse } from "../../../src/apiv2"; import { InstanceType } from "./code-lens-provider"; import { pluginLogger } from "../logger-wrapper"; @@ -63,27 +53,6 @@ export class DataConnectService { return dcs?.getApiServicePathByPath(projectId, path); } - private async decodeResponse( - response: Response, - format?: "application/json", - ): Promise { - const contentType = response.headers.get("Content-Type"); - if (!contentType) { - throw new Error("Invalid content type"); - } - - if (format && !contentType.includes(format)) { - throw new Error( - `Invalid content type. Expected ${format} but got ${contentType}`, - ); - } - - if (contentType.includes("application/json")) { - return response.json(); - } - - return response.text(); - } private async handleProdResponse( response: ClientResponse, ): Promise { @@ -257,32 +226,6 @@ export class DataConnectService { docsLink() { return this.dataConnectToolkit.getGeneratedDocsURL(); } - - // Start cloud section - - async generateOperation( - path: string /** currently unused; instead reading the first service config */, - naturalLanguageQuery: string, - type: "schema" | "operation", - chatHistory: ChatMessage[], - ): Promise { - const client = cloudAICompationClient(); - const servicePath = await this.servicePath( - dataConnectConfigs.value?.tryReadValue?.values[0].path as string, - ); - - if (!servicePath) { - return undefined; - } - - const request: CallCloudAiCompanionRequest = { - servicePath, - naturalLanguageQuery, - chatHistory, - }; - const resp = await callCloudAICompanion(client, request, type); - return resp; - } } function parseVariableString(variables: string): Record { diff --git a/src/dataconnect/cloudAICompanionClient.spec.ts b/src/dataconnect/cloudAICompanionClient.spec.ts deleted file mode 100644 index f7af7125e12..00000000000 --- a/src/dataconnect/cloudAICompanionClient.spec.ts +++ /dev/null @@ -1,94 +0,0 @@ -import { expect } from "chai"; -import * as sinon from "sinon"; -import * as nock from "nock"; -import * as chai from "chai"; -import { callCloudAICompanion, cloudAICompationClient } from "./cloudAICompanionClient"; -import { Client } from "../apiv2"; -import { CallCloudAiCompanionRequest, CloudAICompanionResponse } from "./cloudAICompanionTypes"; - -chai.use(require("chai-as-promised")); - -describe("cloudAICompanionClient", () => { - let sandbox: sinon.SinonSandbox; - - beforeEach(() => { - sandbox = sinon.createSandbox(); - }); - - afterEach(() => { - sandbox.restore(); - nock.cleanAll(); - }); - - describe("callCloudAICompanion", () => { - const fakeRequest: CallCloudAiCompanionRequest = { - servicePath: "projects/my-project/locations/us-central1/services/my-service", - naturalLanguageQuery: "Get all users", - chatHistory: [], - }; - - it("should call the Cloud AI Companion API for schema generation", async () => { - const expectedResponse: CloudAICompanionResponse = { - output: { - messages: [{ author: "MODEL", content: "Generated schema" }], - }, - }; - nock("https://cloudaicompanion.googleapis.com") - .post("/v1/projects/my-project/locations/global/instances/default:completeTask", (body) => { - expect(body.experienceContext.experience).to.equal( - "/appeco/firebase/fdc-schema-generator", - ); - return true; - }) - .reply(200, expectedResponse); - - const client = cloudAICompationClient(); - const response = await callCloudAICompanion(client, fakeRequest, "schema"); - expect(response).to.deep.equal(expectedResponse); - }); - - it("should call the Cloud AI Companion API for operation generation", async () => { - const expectedResponse: CloudAICompanionResponse = { - output: { - messages: [{ author: "MODEL", content: "Generated operation" }], - }, - }; - nock("https://cloudaicompanion.googleapis.com") - .post("/v1/projects/my-project/locations/global/instances/default:completeTask", (body) => { - expect(body.experienceContext.experience).to.equal( - "/appeco/firebase/fdc-query-generator", - ); - return true; - }) - .reply(200, expectedResponse); - - const client = cloudAICompationClient(); - const response = await callCloudAICompanion(client, fakeRequest, "operation"); - expect(response).to.deep.equal(expectedResponse); - }); - - it("should handle errors from the Cloud AI Companion API", async () => { - nock("https://cloudaicompanion.googleapis.com") - .post("/v1/projects/my-project/locations/global/instances/default:completeTask") - .reply(500, { error: { message: "Internal Server Error" } }); - - const client = cloudAICompationClient(); - const response = await callCloudAICompanion(client, fakeRequest, "schema"); - - expect(response.error).to.exist; - expect(response.output.messages).to.deep.equal([]); - }); - - it("should throw an error for an invalid service name", async () => { - const invalidRequest: CallCloudAiCompanionRequest = { - servicePath: "invalid-service-name", - naturalLanguageQuery: "Get all users", - chatHistory: [], - }; - const client = new Client({ urlPrefix: "", apiVersion: "" }); - await expect(callCloudAICompanion(client, invalidRequest, "schema")).to.be.rejectedWith( - "Invalid service name: invalid-service-name", - ); - }); - }); -}); diff --git a/src/dataconnect/cloudAICompanionClient.ts b/src/dataconnect/cloudAICompanionClient.ts deleted file mode 100644 index 1e6544b2b60..00000000000 --- a/src/dataconnect/cloudAICompanionClient.ts +++ /dev/null @@ -1,106 +0,0 @@ -import { Client } from "../apiv2"; -import { cloudAiCompanionOrigin } from "../api"; -import { - CloudAICompanionResponse, - CloudAICompanionRequest, - CloudAICompanionInput, - ClientContext, - CallCloudAiCompanionRequest, -} from "./cloudAICompanionTypes"; -import { FirebaseError } from "../error"; - -const CLOUD_AI_COMPANION_VERSION = "v1"; -const CLIENT_CONTEXT_NAME_IDENTIFIER = "firebase_vscode"; -const FIREBASE_CHAT_REQUEST_CONTEXT_TYPE_NAME = - "type.googleapis.com/google.cloud.cloudaicompanion.v1main.FirebaseChatRequestContext"; -const FDC_SCHEMA_EXPERIENCE_CONTEXT = "/appeco/firebase/fdc-schema-generator"; -const FDC_OPERATION_EXPERIENCE_CONTEXT = "/appeco/firebase/fdc-query-generator"; -const USER_AUTHOR = "USER"; -type GENERATION_TYPE = "schema" | "operation"; - -export function cloudAICompationClient(): Client { - return new Client({ - urlPrefix: cloudAiCompanionOrigin(), - apiVersion: CLOUD_AI_COMPANION_VERSION, - auth: true, - }); -} - -export async function callCloudAICompanion( - client: Client, - vscodeRequest: CallCloudAiCompanionRequest, - type: GENERATION_TYPE, -): Promise { - const request = buildRequest(vscodeRequest, type); - const { projectId } = getServiceParts(vscodeRequest.servicePath); - - const instance = toChatResourceName(projectId); - - try { - const res = await client.post( - `${instance}:completeTask`, - request, - ); - return res.body; - } catch (error: unknown) { - return { output: { messages: [] }, error: error as FirebaseError }; - } -} - -function buildRequest( - { servicePath, naturalLanguageQuery, chatHistory }: CallCloudAiCompanionRequest, - type: GENERATION_TYPE, -): CloudAICompanionRequest { - const { serviceId } = getServiceParts(servicePath); - const input: CloudAICompanionInput = { - messages: [ - ...chatHistory, - { - author: USER_AUTHOR, - content: naturalLanguageQuery, - }, - ], - }; - - const clientContext: ClientContext = { - name: CLIENT_CONTEXT_NAME_IDENTIFIER, - // TODO: determine if we should pass vscode version; // version: ideContext.ver, - additionalContext: { - "@type": FIREBASE_CHAT_REQUEST_CONTEXT_TYPE_NAME, - fdcInfo: { - serviceId, - fdcServiceName: servicePath, - requiresQuery: true, - }, - }, - }; - - return { - input, - clientContext, - experienceContext: { - experience: - type === "schema" ? FDC_SCHEMA_EXPERIENCE_CONTEXT : FDC_OPERATION_EXPERIENCE_CONTEXT, - }, - }; -} - -function toChatResourceName(projectId: string): string { - return `projects/${projectId}/locations/global/instances/default`; -} - -/** Gets service name parts */ -interface ServiceParts { - projectId: string; - locationId: string; - serviceId: string; -} -function getServiceParts(name: string): ServiceParts { - const match = name.match(/projects\/([^/]*)\/locations\/([^/]*)\/services\/([^/]*)/); - - if (!match) { - throw new Error(`Invalid service name: ${name}`); - } - - return { projectId: match[1], locationId: match[2], serviceId: match[3] }; -}