WIP
This commit is contained in:
cghislai
parent
30a44455e8
commit
dd4b116bbb
@ -1,7 +1,13 @@
|
||||
/**
|
||||
* 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';
|
||||
@ -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({
|
||||
model: this.model,
|
||||
tools: geminiProjectProcessor ? this.fileOperationTools : undefined,
|
||||
});
|
||||
|
||||
// Send the AI.md file and additional context to Gemini without hardcoded instructions
|
||||
// Create the prompt
|
||||
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
|
||||
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);
|
||||
|
||||
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
|
||||
@ -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
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user