From dd4b116bbb96cb6f9ab0fa523caec6b85b09c7bd Mon Sep 17 00:00:00 2001 From: cghislai Date: Sun, 8 Jun 2025 04:40:26 +0200 Subject: [PATCH] WIP --- .../src/services/gemini-service.ts | 1018 ++++++++++++----- 1 file changed, 736 insertions(+), 282 deletions(-) diff --git a/src/functions/prompts-to-test-spec/src/services/gemini-service.ts b/src/functions/prompts-to-test-spec/src/services/gemini-service.ts index 792a9ef..44b2e13 100644 --- a/src/functions/prompts-to-test-spec/src/services/gemini-service.ts +++ b/src/functions/prompts-to-test-spec/src/services/gemini-service.ts @@ -1,168 +1,174 @@ /** * Service for handling Gemini API operations */ -import { VertexAI, FunctionDeclaration, Tool, FunctionDeclarationSchemaType } from '@google-cloud/vertexai'; +import { + VertexAI, + FunctionDeclaration, + Tool, + FunctionDeclarationSchemaType, + GenerateContentRequest +} from '@google-cloud/vertexai'; import * as fs from 'fs'; import * as path from 'path'; -import { Project, Workitem } from '../types'; -import { - GOOGLE_CLOUD_PROJECT_ID, - GOOGLE_CLOUD_LOCATION, - GEMINI_MODEL, - DRY_RUN_SKIP_GEMINI, - GOOGLE_API_KEY +import {Project, Workitem} from '../types'; +import { + GOOGLE_CLOUD_PROJECT_ID, + GOOGLE_CLOUD_LOCATION, + GEMINI_MODEL, + DRY_RUN_SKIP_GEMINI, + GOOGLE_API_KEY } from '../config'; export class GeminiService { - private vertexAI: VertexAI; - private model: string; - private projectId: string; - private location: string; - private fileOperationTools: Tool[]; + private vertexAI: VertexAI; + private model: string; + private projectId: string; + private location: string; + private fileOperationTools: Tool[]; - constructor(projectId?: string, location?: string, model?: string) { - this.projectId = projectId || GOOGLE_CLOUD_PROJECT_ID; - this.location = location || GOOGLE_CLOUD_LOCATION; - this.model = model || GEMINI_MODEL; + constructor(projectId?: string, location?: string, model?: string) { + this.projectId = projectId || GOOGLE_CLOUD_PROJECT_ID; + this.location = location || GOOGLE_CLOUD_LOCATION; + this.model = model || GEMINI_MODEL; - if (!this.projectId) { - throw new Error('Google Cloud Project ID is required'); + if (!this.projectId) { + throw new Error('Google Cloud Project ID is required'); + } + + // Initialize VertexAI with default authentication + // Authentication is handled by the Google Auth library, which supports: + // 1. API keys via GOOGLE_API_KEY environment variable + // 2. Service accounts via GOOGLE_APPLICATION_CREDENTIALS environment variable + // 3. Application Default Credentials + // See: https://cloud.google.com/vertex-ai/docs/authentication + this.vertexAI = new VertexAI({ + project: this.projectId, + location: this.location, + }); + + // Define file operation functions + this.fileOperationTools = [ + { + function_declarations: [ + { + name: "getFileContent", + description: "Get the content of a file in the project repository", + parameters: { + type: FunctionDeclarationSchemaType.OBJECT, + properties: { + filePath: { + type: FunctionDeclarationSchemaType.STRING, + description: "Path to the file relative to the project repository root" + } + }, + required: ["filePath"] + } + }, + { + name: "writeFileContent", + description: "Write content to a file in the project repository", + parameters: { + type: FunctionDeclarationSchemaType.OBJECT, + properties: { + filePath: { + type: FunctionDeclarationSchemaType.STRING, + description: "Path to the file relative to the project repository root" + }, + content: { + type: FunctionDeclarationSchemaType.STRING, + description: "Content to write to the file" + } + }, + required: ["filePath", "content"] + } + }, + { + name: "fileExists", + description: "Check if a file exists in the project repository", + parameters: { + type: FunctionDeclarationSchemaType.OBJECT, + properties: { + filePath: { + type: FunctionDeclarationSchemaType.STRING, + description: "Path to the file relative to the project repository root" + } + }, + required: ["filePath"] + } + }, + { + name: "listFiles", + description: "List files in a directory in the project repository", + parameters: { + type: FunctionDeclarationSchemaType.OBJECT, + properties: { + dirPath: { + type: FunctionDeclarationSchemaType.STRING, + description: "Path to the directory relative to the project repository root" + } + }, + required: ["dirPath"] + } + }, + { + name: "grepFiles", + description: "Search for a string in project files", + parameters: { + type: FunctionDeclarationSchemaType.OBJECT, + properties: { + searchString: { + type: FunctionDeclarationSchemaType.STRING, + description: "String to search for in project files" + }, + filePattern: { + type: FunctionDeclarationSchemaType.STRING, + description: "Optional file pattern to limit the search (e.g., '*.ts', 'src/*.java')" + } + }, + required: ["searchString"] + } + }, + { + name: "deleteFile", + description: "Delete a file from the project repository", + parameters: { + type: FunctionDeclarationSchemaType.OBJECT, + properties: { + filePath: { + type: FunctionDeclarationSchemaType.STRING, + description: "Path to the file relative to the project repository root" + } + }, + required: ["filePath"] + } + } + ] + } + ]; } - // Initialize VertexAI with default authentication - // Authentication is handled by the Google Auth library, which supports: - // 1. API keys via GOOGLE_API_KEY environment variable - // 2. Service accounts via GOOGLE_APPLICATION_CREDENTIALS environment variable - // 3. Application Default Credentials - // See: https://cloud.google.com/vertex-ai/docs/authentication - this.vertexAI = new VertexAI({ - project: this.projectId, - location: this.location, - }); + /** + * Generate feature file content using Gemini API with streaming + * @param guidelines Project guidelines + * @param workitemContent Workitem content + * @param workitemName Name of the workitem + * @param geminiProjectProcessor Optional GeminiProjectProcessor to handle function calls + * @param additionalContext Optional additional context from relevant files + * @returns Generated feature file content + */ + async generateFeatureFile( + guidelines: string, + workitemContent: string, + workitemName: string, + geminiProjectProcessor?: any, + additionalContext: string = '' + ): Promise { + const currentDate = new Date().toISOString(); - // Define file operation functions - this.fileOperationTools = [ - { - function_declarations: [ - { - name: "getFileContent", - description: "Get the content of a file in the project repository", - parameters: { - type: FunctionDeclarationSchemaType.OBJECT, - properties: { - filePath: { - type: FunctionDeclarationSchemaType.STRING, - description: "Path to the file relative to the project repository root" - } - }, - required: ["filePath"] - } - }, - { - name: "writeFileContent", - description: "Write content to a file in the project repository", - parameters: { - type: FunctionDeclarationSchemaType.OBJECT, - properties: { - filePath: { - type: FunctionDeclarationSchemaType.STRING, - description: "Path to the file relative to the project repository root" - }, - content: { - type: FunctionDeclarationSchemaType.STRING, - description: "Content to write to the file" - } - }, - required: ["filePath", "content"] - } - }, - { - name: "fileExists", - description: "Check if a file exists in the project repository", - parameters: { - type: FunctionDeclarationSchemaType.OBJECT, - properties: { - filePath: { - type: FunctionDeclarationSchemaType.STRING, - description: "Path to the file relative to the project repository root" - } - }, - required: ["filePath"] - } - }, - { - name: "listFiles", - description: "List files in a directory in the project repository", - parameters: { - type: FunctionDeclarationSchemaType.OBJECT, - properties: { - dirPath: { - type: FunctionDeclarationSchemaType.STRING, - description: "Path to the directory relative to the project repository root" - } - }, - required: ["dirPath"] - } - }, - { - name: "grepFiles", - description: "Search for a string in project files", - parameters: { - type: FunctionDeclarationSchemaType.OBJECT, - properties: { - searchString: { - type: FunctionDeclarationSchemaType.STRING, - description: "String to search for in project files" - }, - filePattern: { - type: FunctionDeclarationSchemaType.STRING, - description: "Optional file pattern to limit the search (e.g., '*.ts', 'src/*.java')" - } - }, - required: ["searchString"] - } - }, - { - name: "deleteFile", - description: "Delete a file from the project repository", - parameters: { - type: FunctionDeclarationSchemaType.OBJECT, - properties: { - filePath: { - type: FunctionDeclarationSchemaType.STRING, - description: "Path to the file relative to the project repository root" - } - }, - required: ["filePath"] - } - } - ] - } - ]; - } - - /** - * Generate feature file content using Gemini API - * @param guidelines Project guidelines - * @param workitemContent Workitem content - * @param workitemName Name of the workitem - * @param geminiProjectProcessor Optional GeminiProjectProcessor to handle function calls - * @param additionalContext Optional additional context from relevant files - * @returns Generated feature file content - */ - async generateFeatureFile( - guidelines: string, - workitemContent: string, - workitemName: string, - geminiProjectProcessor?: any, - additionalContext: string = '' - ): Promise { - const currentDate = new Date().toISOString(); - - // If dry run is enabled, return a mock feature file - if (DRY_RUN_SKIP_GEMINI) { - console.log(`[DRY RUN] Skipping Gemini API call for generating feature file for ${workitemName}`); - return `# Generated by prompts-to-test-spec on ${currentDate} (DRY RUN) + // If dry run is enabled, return a mock feature file + if (DRY_RUN_SKIP_GEMINI) { + console.log(`[DRY RUN] Skipping Gemini API call for generating feature file for ${workitemName}`); + return `# Generated by prompts-to-test-spec on ${currentDate} (DRY RUN) # Source: ${workitemName} Feature: ${workitemName} (DRY RUN) @@ -174,15 +180,10 @@ Feature: ${workitemName} (DRY RUN) When the feature file is generated Then a mock feature file is returned `; - } + } - const generativeModel = this.vertexAI.getGenerativeModel({ - model: this.model, - tools: geminiProjectProcessor ? this.fileOperationTools : undefined, - }); - - // Send the AI.md file and additional context to Gemini without hardcoded instructions - const prompt = ` + // Create the prompt + const prompt = ` ${guidelines} Workitem: @@ -195,7 +196,7 @@ Include the following comment at the top of any generated files: # Generated by prompts-to-test-spec on ${currentDate} # Source: ${workitemName} -You have access to file operations to help you understand the project structure and create better implementations: +You have access to the following function calls to help you understand the project structure and create implementations: - getFileContent(filePath): Get the content of a file in the project repository - writeFileContent(filePath, content): Write content to a file in the project repository - fileExists(filePath): Check if a file exists in the project repository @@ -204,138 +205,591 @@ You have access to file operations to help you understand the project structure - deleteFile(filePath): Delete a file from the project repository You can decide whether to create, update, delete or skip implementing this workitem based on your analysis. -Include the decision in your response. + +In your response, just include your decision with a short motivation in json format. For instance: +{ "decision": "create", "reason": "This workitem was not implemented" } ${additionalContext ? `\nAdditional context from project files:${additionalContext}` : ''} `; - // Start the chat session - const chat = generativeModel.startChat(); + // Initialize Vertex with your Cloud project and location + const vertexAI = new VertexAI({ + project: this.projectId, + location: this.location, + }); - // Send the initial message - const result = await chat.sendMessage(prompt); - console.log(`Gemini result for ${workitemName}`, JSON.stringify(result, null, 2)); + // Instantiate the model with our file operation tools + const generativeModel = vertexAI.getGenerativeModel({ + model: this.model, + tools: geminiProjectProcessor ? this.fileOperationTools : undefined, + generation_config: { + temperature: 0.1, // Very low temperature for more deterministic responses + top_p: 0.95, // Higher top_p to allow more diverse completions when needed + top_k: 40, // Consider only the top 40 tokens + }, + }); + // Create the initial request + const request: GenerateContentRequest = { + contents: [ + {role: 'user', parts: [{text: prompt}]} + ], + tools: geminiProjectProcessor ? this.fileOperationTools : undefined, + }; - // Process function calls if needed - let finalResponse = await this.processFunctionCalls(result, chat, geminiProjectProcessor); - console.log(`Gemini response for ${workitemName}: ${finalResponse}`, result.response.candidates[0]); + // Generate content in a streaming fashion + const streamingResp = await generativeModel.generateContentStream(request); - return finalResponse; - } + let finalResponse = ''; + let pendingFunctionCalls = []; - /** - * Process function calls in the Gemini response - * @param result The result from Gemini - * @param chat The chat session - * @param geminiProjectProcessor The GeminiProjectProcessor to handle function calls - * @returns The final generated text - */ - private async processFunctionCalls(result: any, chat: any, geminiProjectProcessor?: any): Promise { - // Check if there are function calls in the response - if (!result.functionCalls || result.functionCalls.length === 0 || !geminiProjectProcessor) { - // No function calls, return the text response - // Access text content from the response structure in @google-cloud/vertexai v0.5.0 - return result.candidates?.[0]?.content?.parts?.[0]?.text || ''; - } + // Process the streaming response + for await (const item of streamingResp.stream) { + // Check if there's a function call in any part of the response + let functionCall = null; + let textContent = ''; - console.log(`Processing ${result.functionCalls.length} function calls from Gemini`); + // Iterate over every part in the response + for (const part of item.candidates?.[0]?.content?.parts || []) { + if (part.functionCall) { + functionCall = part.functionCall; + break; + } else if (part.text) { + textContent += part.text; + } + } - // Process each function call - for (const functionCall of result.functionCalls) { - const functionName = functionCall.name; - const functionArgs = JSON.parse(functionCall.args); - - console.log(`Executing function: ${functionName} with args:`, functionArgs); - - let functionResponse; - try { - // Execute the function using the GeminiProjectProcessor - switch (functionName) { - case 'getFileContent': - functionResponse = geminiProjectProcessor.getFileContent(functionArgs.filePath); - break; - case 'writeFileContent': - // Get the current workitem name from the context - const currentWorkitem = geminiProjectProcessor.getCurrentWorkitem(); - geminiProjectProcessor.writeFileContent(functionArgs.filePath, functionArgs.content, currentWorkitem?.name); - functionResponse = `File ${functionArgs.filePath} written successfully`; - break; - case 'fileExists': - functionResponse = geminiProjectProcessor.fileExists(functionArgs.filePath); - break; - case 'listFiles': - functionResponse = geminiProjectProcessor.listFiles(functionArgs.dirPath); - break; - case 'grepFiles': - functionResponse = geminiProjectProcessor.grepFiles(functionArgs.searchString, functionArgs.filePattern); - break; - case 'deleteFile': - functionResponse = geminiProjectProcessor.deleteFile(functionArgs.filePath); - break; - default: - throw new Error(`Unknown function: ${functionName}`); + if (functionCall) { + console.log(`Function call detected: ${functionCall.name}`); + pendingFunctionCalls.push(functionCall); + } else if (textContent) { + // If there's text, append it to the final response + finalResponse += textContent; + } } - // Send the function response back to Gemini - const functionResponseObj = { functionResponse: { name: functionName, response: { result: JSON.stringify(functionResponse) } } }; - const nextResult = await chat.sendMessage(functionResponseObj); + // Process any function calls that were detected + if (pendingFunctionCalls.length > 0 && geminiProjectProcessor) { + console.log(`Processing ${pendingFunctionCalls.length} function calls from streaming response`); - // Recursively process any additional function calls - return this.processFunctionCalls(nextResult, chat, geminiProjectProcessor); - } catch (error) { - console.error(`Error executing function ${functionName}:`, error); + let currentRequest: GenerateContentRequest = request; - // Send the error back to Gemini - const errorResponseObj = { - functionResponse: { - name: functionName, - response: { error: error instanceof Error ? error.message : String(error) } - } + // Process each function call + for (const functionCall of pendingFunctionCalls) { + const functionName = functionCall.name; + const functionArgs = (typeof functionCall.args === 'string' ? + JSON.parse(functionCall.args) : functionCall.args) as { + filePath?: string; + content?: string; + dirPath?: string; + searchString?: string; + filePattern?: string; + }; + + console.log(`Executing function: ${functionName} with args:`, functionArgs); + + let functionResponse; + try { + // Execute the function using the GeminiProjectProcessor + switch (functionName) { + case 'getFileContent': + functionResponse = geminiProjectProcessor.getFileContent(functionArgs.filePath!); + break; + case 'writeFileContent': + // Get the current workitem name from the context + const currentWorkitem = geminiProjectProcessor.getCurrentWorkitem(); + geminiProjectProcessor.writeFileContent(functionArgs.filePath!, functionArgs.content!, currentWorkitem?.name); + functionResponse = `File ${functionArgs.filePath} written successfully`; + break; + case 'fileExists': + functionResponse = geminiProjectProcessor.fileExists(functionArgs.filePath!); + break; + case 'listFiles': + functionResponse = geminiProjectProcessor.listFiles(functionArgs.dirPath!); + break; + case 'grepFiles': + functionResponse = geminiProjectProcessor.grepFiles(functionArgs.searchString!, functionArgs.filePattern); + break; + case 'deleteFile': + functionResponse = geminiProjectProcessor.deleteFile(functionArgs.filePath!); + break; + default: + throw new Error(`Unknown function: ${functionName}`); + } + + // Create a function response object + const functionResponseObj = { + name: functionName, + response: {result: JSON.stringify(functionResponse)} + }; + + // Update the request with the function call and response + currentRequest = { + contents: [ + ...currentRequest.contents, + { + role: 'ASSISTANT', + parts: [ + { + functionCall: functionCall + } + ] + }, + { + role: 'USER', + parts: [ + { + functionResponse: functionResponseObj + } + ] + } + ], + tools: geminiProjectProcessor ? this.fileOperationTools : undefined, + }; + + // Generate the next response + const nextStreamingResp = await generativeModel.generateContentStream(currentRequest); + + // Process the next streaming response + for await (const nextItem of nextStreamingResp.stream) { + let textContent = ''; + + // Iterate over every part in the response + for (const part of nextItem.candidates?.[0]?.content?.parts || []) { + if (part.text) { + textContent += part.text; + } + } + + if (textContent) { + finalResponse += textContent; + } + } + + } catch (error) { + console.error(`Error executing function ${functionName}:`, error); + + // Create an error response object + const errorResponseObj = { + name: functionName, + response: {error: error instanceof Error ? error.message : String(error)} + }; + + // Update the request with the function call and error response + currentRequest = { + contents: [ + ...currentRequest.contents, + { + role: 'ASSISTANT', + parts: [ + { + functionCall: functionCall + } + ] + }, + { + role: 'USER', + parts: [ + { + functionResponse: errorResponseObj + } + ] + } + ], + tools: geminiProjectProcessor ? this.fileOperationTools : undefined, + }; + + // Generate the next response + const nextStreamingResp = await generativeModel.generateContentStream(currentRequest); + + // Process the next streaming response + for await (const nextItem of nextStreamingResp.stream) { + let textContent = ''; + + // Iterate over every part in the response + for (const part of nextItem.candidates?.[0]?.content?.parts || []) { + if (part.text) { + textContent += part.text; + } + } + + if (textContent) { + finalResponse += textContent; + } + } + } + } + } + + console.log(`Gemini response for ${workitemName}: ${finalResponse}`); + return finalResponse; + } + + /** + * Process function calls in the Gemini response + * @param result The result from Gemini + * @param chat The chat session + * @param geminiProjectProcessor The GeminiProjectProcessor to handle function calls + * @returns The final generated text + */ + private async processFunctionCalls(result: any, chat: any, geminiProjectProcessor?: any): Promise { + // Check if there are function calls in the response + // Function calls can be at the top level or nested within candidates + const functionCalls = result.functionCalls || + (result.response?.candidates?.[0]?.functionCall ? [result.response.candidates[0].functionCall] : []) || + (result.candidates?.[0]?.functionCall ? [result.candidates[0].functionCall] : []); + + if (functionCalls.length === 0 || !geminiProjectProcessor) { + // No function calls, return the text response + // Access text content from the response structure in @google-cloud/vertexai v0.5.0 + return result.candidates?.[0]?.content?.parts?.[0]?.text || + result.response?.candidates?.[0]?.content?.parts?.[0]?.text || ''; + } + + console.log(`Processing ${functionCalls.length} function calls from Gemini`); + + // Process each function call + for (const functionCall of functionCalls) { + const functionName = functionCall.name; + // Handle both cases: when args is already an object and when it's a string that needs to be parsed + const functionArgs = (typeof functionCall.args === 'string' ? + JSON.parse(functionCall.args) : functionCall.args) as { + filePath?: string; + content?: string; + dirPath?: string; + searchString?: string; + filePattern?: string; + }; + + console.log(`Executing function: ${functionName} with args:`, functionArgs); + + let functionResponse; + try { + // Execute the function using the GeminiProjectProcessor + switch (functionName) { + case 'getFileContent': + functionResponse = geminiProjectProcessor.getFileContent(functionArgs.filePath!); + break; + case 'writeFileContent': + // Get the current workitem name from the context + const currentWorkitem = geminiProjectProcessor.getCurrentWorkitem(); + geminiProjectProcessor.writeFileContent(functionArgs.filePath!, functionArgs.content!, currentWorkitem?.name); + functionResponse = `File ${functionArgs.filePath} written successfully`; + break; + case 'fileExists': + functionResponse = geminiProjectProcessor.fileExists(functionArgs.filePath!); + break; + case 'listFiles': + functionResponse = geminiProjectProcessor.listFiles(functionArgs.dirPath!); + break; + case 'grepFiles': + functionResponse = geminiProjectProcessor.grepFiles(functionArgs.searchString!, functionArgs.filePattern); + break; + case 'deleteFile': + functionResponse = geminiProjectProcessor.deleteFile(functionArgs.filePath!); + break; + default: + throw new Error(`Unknown function: ${functionName}`); + } + + // Send the function response back to Gemini + const functionResponseObj = { + functionResponse: { + name: functionName, + response: {result: JSON.stringify(functionResponse)} + } + }; + const nextResult = await chat.sendMessage(functionResponseObj); + + // Recursively process any additional function calls + return this.processFunctionCalls(nextResult, chat, geminiProjectProcessor); + } catch (error) { + console.error(`Error executing function ${functionName}:`, error); + + // Send the error back to Gemini + const errorResponseObj = { + functionResponse: { + name: functionName, + response: {error: error instanceof Error ? error.message : String(error)} + } + }; + const nextResult = await chat.sendMessage(errorResponseObj); + + // Recursively process any additional function calls + return this.processFunctionCalls(nextResult, chat, geminiProjectProcessor); + } + } + + // Return the final text response + // Access text content from the response structure in @google-cloud/vertexai v0.5.0 + return result.candidates?.[0]?.content?.parts?.[0]?.text || ''; + } + + /** + * Example of using function calling with streaming content generation + * This method demonstrates how to use the Vertex AI API for function calling in a streaming context + * @param projectId Google Cloud project ID + * @param location Google Cloud location + * @param model Gemini model to use + */ + async functionCallingStreamExample( + projectId: string = this.projectId, + location: string = this.location, + model: string = this.model + ): Promise { + // Initialize Vertex with your Cloud project and location + const vertexAI = new VertexAI({project: projectId, location: location}); + + // Instantiate the model + const generativeModel = vertexAI.getGenerativeModel({ + model: model, + }); + + // Define the function declaration for the weather function + const functionDeclarations: FunctionDeclaration[] = [ + { + name: "get_current_weather", + description: "Get the current weather in a given location", + parameters: { + type: FunctionDeclarationSchemaType.OBJECT, + properties: { + location: { + type: FunctionDeclarationSchemaType.STRING, + description: "The city and state, e.g., San Francisco, CA" + } + }, + required: ["location"] + } + } + ]; + + // Create a mock function response + const functionResponseParts = [ + { + functionResponse: { + name: "get_current_weather", + response: { + temperature: "72", + unit: "fahrenheit", + description: "Sunny" + } + } + } + ]; + + // Create the request with function calling + const request = { + contents: [ + {role: 'user', parts: [{text: 'What is the weather in Boston?'}]}, + { + role: 'ASSISTANT', + parts: [ + { + functionCall: { + name: 'get_current_weather', + args: {location: 'Boston'}, + }, + }, + ], + }, + {role: 'USER', parts: functionResponseParts}, + ], + tools: [{function_declarations: functionDeclarations}], }; - const nextResult = await chat.sendMessage(errorResponseObj); - // Recursively process any additional function calls - return this.processFunctionCalls(nextResult, chat, geminiProjectProcessor); - } + // Generate content in a streaming fashion + const streamingResp = await generativeModel.generateContentStream(request); + + // Process the streaming response + for await (const item of streamingResp.stream) { + // Iterate over every part in the response + for (const part of item.candidates?.[0]?.content?.parts || []) { + if (part.text) { + console.log(part.text); + } + } + } } - // Return the final text response - // Access text content from the response structure in @google-cloud/vertexai v0.5.0 - return result.candidates?.[0]?.content?.parts?.[0]?.text || ''; - } + /** + * Example of using function calling with streaming content generation for file operations + * This method demonstrates how to use the Vertex AI API for file operation function calling in a streaming context + * @param workitemName Name of the workitem + * @param geminiProjectProcessor The GeminiProjectProcessor to handle function calls + */ + async fileOperationsStreamExample( + workitemName: string, + geminiProjectProcessor: any + ): Promise { + // Initialize Vertex with your Cloud project and location + const vertexAI = new VertexAI({ + project: this.projectId, + location: this.location, + }); - /** - * Generate a pull request description using Gemini API - * @param processedWorkitems List of processed workitems - * @param gitPatch Optional git patch showing code changes - * @returns Generated pull request description - */ - async generatePullRequestDescription( - processedWorkitems: { workitem: Workitem; success: boolean; error?: string }[], - gitPatch?: string - ): Promise { - // Prepare workitem data for the prompt - const all: string[] = []; + // Instantiate the model with our file operation tools + const generativeModel = vertexAI.getGenerativeModel({ + model: this.model, + tools: this.fileOperationTools, + }); - for (const item of processedWorkitems) { - const { workitem, success, error } = item; + // Create a prompt that asks the model to check if a file exists and create it if it doesn't + const prompt = `Check if the file 'example.txt' exists and create it with some content if it doesn't.`; - // Add all workitems to the all list regardless of status - all.push(`- ${workitem.name}`); + // Create the initial request + const request = { + contents: [ + {role: 'user', parts: [{text: prompt}]} + ], + tools: this.fileOperationTools, + }; + + // Generate content in a streaming fashion + const streamingResp = await generativeModel.generateContentStream(request); + + // Process the streaming response + for await (const item of streamingResp.stream) { + // Check if there's a function call in any part of the response + let functionCall = null; + + // Iterate over every part in the response + for (const part of item.candidates?.[0]?.content?.parts || []) { + if (part.functionCall) { + functionCall = part.functionCall; + break; + } + } + + if (functionCall) { + console.log(`Function call detected: ${functionCall.name}`); + + // Execute the function + const functionName = functionCall.name; + const functionArgs = functionCall.args as { + filePath?: string; + content?: string; + dirPath?: string; + searchString?: string; + filePattern?: string; + }; + + console.log(`Executing function: ${functionName} with args:`, functionArgs); + + let functionResponse; + try { + // Execute the function using the GeminiProjectProcessor + switch (functionName) { + case 'getFileContent': + functionResponse = geminiProjectProcessor.getFileContent(functionArgs.filePath!); + break; + case 'writeFileContent': + geminiProjectProcessor.writeFileContent(functionArgs.filePath!, functionArgs.content!, workitemName); + functionResponse = `File ${functionArgs.filePath} written successfully`; + break; + case 'fileExists': + functionResponse = geminiProjectProcessor.fileExists(functionArgs.filePath!); + break; + case 'listFiles': + functionResponse = geminiProjectProcessor.listFiles(functionArgs.dirPath!); + break; + case 'grepFiles': + functionResponse = geminiProjectProcessor.grepFiles(functionArgs.searchString!, functionArgs.filePattern); + break; + case 'deleteFile': + functionResponse = geminiProjectProcessor.deleteFile(functionArgs.filePath!); + break; + default: + throw new Error(`Unknown function: ${functionName}`); + } + + // Create a new request with the function response + const functionResponseObj = { + name: functionName, + response: {result: JSON.stringify(functionResponse)} + }; + + // Continue the conversation with the function response + const nextRequest = { + contents: [ + {role: 'user', parts: [{text: prompt}]}, + { + role: 'ASSISTANT', + parts: [ + { + functionCall: functionCall + } + ] + }, + { + role: 'USER', + parts: [ + { + functionResponse: functionResponseObj + } + ] + } + ], + tools: this.fileOperationTools, + }; + + // Generate the next response + const nextStreamingResp = await generativeModel.generateContentStream(nextRequest); + + // Process the next streaming response + for await (const nextItem of nextStreamingResp.stream) { + // Iterate over every part in the response + for (const part of nextItem.candidates?.[0]?.content?.parts || []) { + if (part.text) { + console.log(part.text); + } + } + } + + } catch (error) { + console.error(`Error executing function ${functionName}:`, error); + } + } else { + // If there's no function call, just log the text from all parts + for (const part of item.candidates?.[0]?.content?.parts || []) { + if (part.text) { + console.log(part.text); + } + } + } + } } - // Create a structured summary of changes - let workitemSummary = ''; + /** + * Generate a pull request description using Gemini API + * @param processedWorkitems List of processed workitems + * @param gitPatch Optional git patch showing code changes + * @returns Generated pull request description + */ + async generatePullRequestDescription( + processedWorkitems: { workitem: Workitem; success: boolean; error?: string }[], + gitPatch?: string + ): Promise { + // Prepare workitem data for the prompt + const all: string[] = []; - if (all.length > 0) { - workitemSummary += 'All workitems:\n' + all.join('\n') + '\n\n'; - } + for (const item of processedWorkitems) { + const {workitem, success, error} = item; - // If dry run is enabled, return a mock PR description - if (DRY_RUN_SKIP_GEMINI) { - console.log(`[DRY RUN] Skipping Gemini API call for generating pull request description`); - return `# Automated PR: Update Workitems (DRY RUN) + // Add all workitems to the all list regardless of status + all.push(`- ${workitem.name}`); + } + + // Create a structured summary of changes + let workitemSummary = ''; + + if (all.length > 0) { + workitemSummary += 'All workitems:\n' + all.join('\n') + '\n\n'; + } + + // If dry run is enabled, return a mock PR description + if (DRY_RUN_SKIP_GEMINI) { + console.log(`[DRY RUN] Skipping Gemini API call for generating pull request description`); + return `# Automated PR: Update Workitems (DRY RUN) This pull request was automatically generated by the prompts-to-test-spec function in dry run mode. @@ -344,16 +798,16 @@ This pull request was automatically generated by the prompts-to-test-spec functi ${workitemSummary} *Note: This is a mock PR description generated during dry run. No actual Gemini API call was made.*`; - } + } - const generativeModel = this.vertexAI.getGenerativeModel({ - model: this.model, - }); + const generativeModel = this.vertexAI.getGenerativeModel({ + model: this.model, + }); - // Prepare the git patch section if available - let gitPatchSection = ''; - if (gitPatch && gitPatch !== "No changes detected.") { - gitPatchSection = ` + // Prepare the git patch section if available + let gitPatchSection = ''; + if (gitPatch && gitPatch !== "No changes detected.") { + gitPatchSection = ` ## Code Changes The following code changes were made in this pull request: @@ -361,9 +815,9 @@ The following code changes were made in this pull request: ${gitPatch} \`\`\` `; - } + } - const prompt = ` + const prompt = ` You are tasked with creating a pull request description for changes to test specifications. The following is a summary of the changes made: @@ -380,11 +834,11 @@ Create a clear, professional pull request description that: The pull request description should be ready to use without further editing. `; - const result = await generativeModel.generateContent(prompt); + const result = await generativeModel.generateContent(prompt); - const response = await result.response; - const generatedText = response.candidates[0]?.content?.parts[0]?.text || ''; + const response = await result.response; + const generatedText = response.candidates[0]?.content?.parts[0]?.text || ''; - return generatedText; - } + return generatedText; + } }