WIP

src/functions/prompts-to-test-spec/src/services/gemini-service.ts

@@ -1,10 +1,16 @@
 /**
  * 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 {Project, Workitem} from '../types';
 import {
     GOOGLE_CLOUD_PROJECT_ID,
     GOOGLE_CLOUD_LOCATION,
@@ -142,7 +148,7 @@ export class GeminiService {
     }
 
     /**
-   * Generate feature file content using Gemini API
+     * Generate feature file content using Gemini API with streaming
      * @param guidelines Project guidelines
      * @param workitemContent Workitem content
      * @param workitemName Name of the workitem
@@ -176,12 +182,7 @@ Feature: ${workitemName} (DRY RUN)
 `;
         }
 
-    const generativeModel = this.vertexAI.getGenerativeModel({
+        // Create the prompt
-      model: this.model,
-      tools: geminiProjectProcessor ? this.fileOperationTools : undefined,
-    });
-
-    // Send the AI.md file and additional context to Gemini without hardcoded instructions
         const prompt = `
 ${guidelines}
 
@@ -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,23 +205,223 @@ 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
+        // Initialize Vertex with your Cloud project and location
-    const chat = generativeModel.startChat();
+        const vertexAI = new VertexAI({
+            project: this.projectId,
+            location: this.location,
+        });
 
-    // Send the initial message
+        // Instantiate the model with our file operation tools
-    const result = await chat.sendMessage(prompt);
+        const generativeModel = vertexAI.getGenerativeModel({
-    console.log(`Gemini result for ${workitemName}`, JSON.stringify(result, null, 2));
+            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
+        // Generate content in a streaming fashion
-    let finalResponse = await this.processFunctionCalls(result, chat, geminiProjectProcessor);
+        const streamingResp = await generativeModel.generateContentStream(request);
-    console.log(`Gemini response for ${workitemName}: ${finalResponse}`,  result.response.candidates[0]);
 
+        let finalResponse = '';
+        let pendingFunctionCalls = [];
+
+        // 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 = '';
+
+            // 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;
+                }
+            }
+
+            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;
+            }
+        }
+
+        // Process any function calls that were detected
+        if (pendingFunctionCalls.length > 0 && geminiProjectProcessor) {
+            console.log(`Processing ${pendingFunctionCalls.length} function calls from streaming response`);
+
+            let currentRequest: GenerateContentRequest = request;
+
+            // 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;
     }
 
@@ -233,18 +434,32 @@ ${additionalContext ? `\nAdditional context from project files:${additionalConte
      */
     private async processFunctionCalls(result: any, chat: any, geminiProjectProcessor?: any): Promise<string> {
         // Check if there are function calls in the response
-    if (!result.functionCalls || result.functionCalls.length === 0 || !geminiProjectProcessor) {
+        // 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 || '';
+            return result.candidates?.[0]?.content?.parts?.[0]?.text ||
+                result.response?.candidates?.[0]?.content?.parts?.[0]?.text || '';
         }
 
-    console.log(`Processing ${result.functionCalls.length} function calls from Gemini`);
+        console.log(`Processing ${functionCalls.length} function calls from Gemini`);
 
         // Process each function call
-    for (const functionCall of result.functionCalls) {
+        for (const functionCall of functionCalls) {
             const functionName = functionCall.name;
-      const functionArgs = JSON.parse(functionCall.args);
+            // 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);
 
@@ -253,32 +468,37 @@ ${additionalContext ? `\nAdditional context from project files:${additionalConte
                 // Execute the function using the GeminiProjectProcessor
                 switch (functionName) {
                     case 'getFileContent':
-            functionResponse = geminiProjectProcessor.getFileContent(functionArgs.filePath);
+                        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);
+                        geminiProjectProcessor.writeFileContent(functionArgs.filePath!, functionArgs.content!, currentWorkitem?.name);
                         functionResponse = `File ${functionArgs.filePath} written successfully`;
                         break;
                     case 'fileExists':
-            functionResponse = geminiProjectProcessor.fileExists(functionArgs.filePath);
+                        functionResponse = geminiProjectProcessor.fileExists(functionArgs.filePath!);
                         break;
                     case 'listFiles':
-            functionResponse = geminiProjectProcessor.listFiles(functionArgs.dirPath);
+                        functionResponse = geminiProjectProcessor.listFiles(functionArgs.dirPath!);
                         break;
                     case 'grepFiles':
-            functionResponse = geminiProjectProcessor.grepFiles(functionArgs.searchString, functionArgs.filePattern);
+                        functionResponse = geminiProjectProcessor.grepFiles(functionArgs.searchString!, functionArgs.filePattern);
                         break;
                     case 'deleteFile':
-            functionResponse = geminiProjectProcessor.deleteFile(functionArgs.filePath);
+                        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 functionResponseObj = {
+                    functionResponse: {
+                        name: functionName,
+                        response: {result: JSON.stringify(functionResponse)}
+                    }
+                };
                 const nextResult = await chat.sendMessage(functionResponseObj);
 
                 // Recursively process any additional function calls
@@ -290,7 +510,7 @@ ${additionalContext ? `\nAdditional context from project files:${additionalConte
                 const errorResponseObj = {
                     functionResponse: {
                         name: functionName,
-            response: { error: error instanceof Error ? error.message : String(error) } 
+                        response: {error: error instanceof Error ? error.message : String(error)}
                     }
                 };
                 const nextResult = await chat.sendMessage(errorResponseObj);
@@ -305,6 +525,240 @@ ${additionalContext ? `\nAdditional context from project files:${additionalConte
         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<void> {
+        // 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}],
+        };
+
+        // 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);
+                }
+            }
+        }
+    }
+
+    /**
+     * 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<void> {
+        // Initialize Vertex with your Cloud project and location
+        const vertexAI = new VertexAI({
+            project: this.projectId,
+            location: this.location,
+        });
+
+        // Instantiate the model with our file operation tools
+        const generativeModel = vertexAI.getGenerativeModel({
+            model: this.model,
+            tools: this.fileOperationTools,
+        });
+
+        // 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.`;
+
+        // 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);
+                    }
+                }
+            }
+        }
+    }
+
     /**
      * Generate a pull request description using Gemini API
      * @param processedWorkitems List of processed workitems
@@ -319,7 +773,7 @@ ${additionalContext ? `\nAdditional context from project files:${additionalConte
         const all: string[] = [];
 
         for (const item of processedWorkitems) {
-      const { workitem, success, error } = item;
+            const {workitem, success, error} = item;
 
             // Add all workitems to the all list regardless of status
             all.push(`- ${workitem.name}`);