genaicode 0.0.33 → 0.0.34

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "genaicode",
-  "version": "0.0.33",
+  "version": "0.0.34",
   "author": "Grzegorz Tańczyk",
   "repository": {
     "type": "git",
@@ -45,11 +45,12 @@
   "dependencies": {
     "@anthropic-ai/sdk": "^0.24.3",
     "@anthropic-ai/vertex-sdk": "^0.4.1",
-    "@google-cloud/vertexai": "^1.3.0",
     "@google-cloud/aiplatform": "^3.25.0",
+    "@google-cloud/vertexai": "^1.3.0",
     "@imgly/background-removal-node": "^1.4.5",
     "diff": "^5.2.0",
     "image-size": "^1.1.1",
+    "jsonschema": "^1.4.1",
     "mime-types": "^2.1.35",
     "openai": "^4.52.7",
     "please-upgrade-node": "^3.2.0",

package/src/ai-service/anthropic.js

@@ -22,6 +22,7 @@ export async function generateContent(prompt, functionDefs, requiredFunctionName
               tool_use_id: response.call_id ?? response.name,
               content: response.content,
               type: 'tool_result',
+              is_error: response.isError === true,
             })),
             ...(item.images ?? []).map((image) => ({
               type: 'image',

package/src/ai-service/common.js

@@ -1,5 +1,5 @@
 import assert from 'node:assert';
-import { functionDefs } from './function-calling.js';
+import { functionDefs } from '../prompt/function-calling.js';
 
 /**
  * Common function to print token usage and estimated cost

package/src/ai-service/common.test.js

@@ -1,6 +1,6 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { printTokenUsageAndCost, processFunctionCalls } from './common.js';
-import { functionDefs } from './function-calling.js';
+import { functionDefs } from '../prompt/function-calling.js';
 
 // Mock cli-params.js
 vi.mock('../cli/cli-params.js', () => ({
@@ -12,6 +12,7 @@ vi.mock('../cli/cli-params.js', () => ({
   cheap: false,
   geminiBlockNone: false,
   requireExplanations: false,
+  vision: false,
 }));
 
 describe('printTokenUsageAndCost', () => {

package/src/ai-service/vertex-ai-claude.js

@@ -28,6 +28,7 @@ export async function generateContent(prompt, functionDefs, requiredFunctionName
               type: 'tool_result',
               tool_use_id: response.call_id ?? response.name,
               content: response.content,
+              is_error: response.isError === true,
             })),
             ...(item.images ?? []).map((image) => ({
               type: 'image',

package/src/main/codegen.test.js

@@ -209,7 +209,7 @@ describe('runCodegen', () => {
       {
         name: 'codegenSummary',
         args: {
-          files: [{ path: 'test.js', updateToolName: 'updateFile' }],
+          fileUpdates: [{ path: 'test.js', updateToolName: 'updateFile' }],
           contextPaths: [],
           explanation: 'Mock summary with image generation failure',
         },
@@ -237,7 +237,7 @@ describe('runCodegen', () => {
       {
         name: 'codegenSummary',
         args: {
-          files: [{ path: 'test.js', updateToolName: 'updateFile' }],
+          fileUpdates: [{ path: 'test.js', updateToolName: 'updateFile' }],
           contextPaths: [],
           explanation: 'Mock summary with image generation failure',
         },
@@ -293,7 +293,7 @@ describe('runCodegen', () => {
       {
         name: 'codegenSummary',
         args: {
-          files: [{ path: 'test.js', updateToolName: 'updateFile' }],
+          fileUpdates: [{ path: 'test.js', updateToolName: 'updateFile' }],
           contextPaths: [],
           explanation: 'Mock summary with cheap image generation',
         },

package/src/prompt/function-calling-validate.js

@@ -0,0 +1,25 @@
+import { Validator } from 'jsonschema';
+import { functionDefs } from './function-calling.js';
+
+export function validateFunctionCall(call) {
+  const validator = new Validator();
+  const functionDef = functionDefs.find((def) => def.name === call.name);
+
+  if (!functionDef) {
+    return {
+      errors: [
+        {
+          message: `Function "${call.name}" is not defined.`,
+        },
+      ],
+    };
+  }
+
+  const validationResult = validator.validate(call.args, functionDef.parameters);
+
+  if (validationResult.errors.length > 0) {
+    return validationResult;
+  }
+
+  return undefined;
+}

package/src/{ai-service → prompt}/function-calling.js

@@ -1,4 +1,4 @@
-import { requireExplanations, temperature, cheap } from '../cli/cli-params.js';
+import { requireExplanations, temperature, cheap, vision } from '../cli/cli-params.js';
 
 /**
  * Function definitions for function calling feature
@@ -43,24 +43,38 @@ export const functionDefs = [
   {
     name: 'codegenSummary',
     description:
-      'This function is called with a summary of proposed updates. It contains a list of file paths that will be subject to code generation request, and also a list of file paths that make sense to use as a context for code generation requests.',
+      'This function is called with a summary of proposed updates.' +
+      'Summary is a data structure(object) which contains:\n' +
+      '- `fileUpdates`: list of proposed file updates that will be subject to subsequent code generation request\n' +
+      '- `contextPaths`: list of file paths that make sense to use as a context for code generation requests.' +
+      '- `explanation`: general explanation of planned code generation updates\n' +
+      'It is crticially important to adhere to the schema of parameters',
     parameters: {
       type: 'object',
       properties: {
-        files: {
+        contextPaths: {
           type: 'array',
-          description: 'An array of proposed file updates.',
+          description:
+            'An array of absolute paths of files that should be used to provide context for the following updates. Context files could be for example the dependencies, or files that depend on one of the files that we want to update in the next step.',
+          items: {
+            type: 'string',
+          },
+        },
+        fileUpdates: {
+          type: 'array',
+          description: 'An array of proposed file updates, each update is an object with several properties.',
           items: {
             type: 'object',
-            description: 'Proposed update of a file, the path, and the method of update',
+            description:
+              'Proposed update of a file. The update is an object which contains properties like file path, update tool name, and few other important properties.',
             properties: {
               path: { type: 'string', description: 'An absolute path of the project file that will be updated' },
               updateToolName: {
                 type: 'string',
                 enum: [
+                  'createFile',
                   'updateFile',
                   'patchFile',
-                  'createFile',
                   'deleteFile',
                   'createDirectory',
                   'moveFile',
@@ -92,19 +106,12 @@ export const functionDefs = [
               contextImageAssets: {
                 type: 'array',
                 description:
-                  'A list of of absolute image asset paths that should be included to the context of LLm request',
+                  'A list of of absolute image asset paths that should be included to the context of LLM request.' +
+                  'This parameter must be used if there is a need to analyze an image.',
                 items: { type: 'string' },
               },
             },
-            required: ['path', 'updateToolName', 'temperature', 'prompt', 'contextImageAssets', 'cheap'],
-          },
-        },
-        contextPaths: {
-          type: 'array',
-          description:
-            'An array of absolute paths of files that should be used to provide context for the following updates. Context files could be for example the dependencies, or files that depend on one of the files that we want to update in the next step.',
-          items: {
-            type: 'string',
+            required: ['path', 'updateToolName', ...(vision ? ['contextImageAssets'] : [])],
           },
         },
         explanation: {
@@ -112,7 +119,7 @@ export const functionDefs = [
           description: 'Explanation of planned changes or explanation of reasoning for no code changes',
         },
       },
-      required: ['files', 'contextPaths', 'explanation'],
+      required: ['fileUpdates', 'contextPaths'],
     },
   },
   {
@@ -183,11 +190,11 @@ Index: filename.js
       properties: {
         filePath: {
           type: 'string',
-          description: 'The file path to create.',
+          description: 'Path of the file that will be created, it must not be empty.',
         },
         newContent: {
           type: 'string',
-          description: 'The content for the new file.',
+          description: 'Content of the file that will be created, it must no be empty.',
         },
         explanation: {
           type: 'string',

package/src/prompt/prompt-service.js

@@ -5,9 +5,10 @@ import mime from 'mime-types';
 
 import { getSystemPrompt } from './systemprompt.js';
 import { getCodeGenPrompt } from './prompt-codegen.js';
-import { functionDefs } from '../ai-service/function-calling.js';
+import { functionDefs } from './function-calling.js';
 import { getSourceCode, getImageAssets } from '../files/read-files.js';
 import { disableContextOptimization, temperature, vision, cheap } from '../cli/cli-params.js';
+import { validateFunctionCall } from './function-calling-validate.js';
 
 /** A function that communicates with model using */
 export async function promptService(generateContentFn, generateImageFn, codegenPrompt = getCodeGenPrompt()) {
@@ -43,16 +44,20 @@ export async function promptService(generateContentFn, generateImageFn, codegenP
 
   prompt.slice(-1)[0].text = messages.prompt;
 
-  let baseResult = await generateContentFn(prompt, functionDefs, 'codegenSummary', temperature, cheap);
+  let baseRequest = [prompt, functionDefs, 'codegenSummary', temperature, cheap];
+  let baseResult = await generateContentFn(...baseRequest);
 
-  const codegenSummaryRequest = baseResult.find((call) => call.name === 'codegenSummary');
+  let codegenSummaryRequest = baseResult.find((call) => call.name === 'codegenSummary');
 
   if (codegenSummaryRequest) {
     // Second stage: for each file request the actual code updates
     console.log('Received codegen summary, will collect partial updates', codegenSummaryRequest.args);
 
+    baseResult = await validateAndRecoverSingleResult(baseRequest, baseResult, messages, generateContentFn);
+    codegenSummaryRequest = baseResult.find((call) => call.name === 'codegenSummary');
+
     // Sometimes the result happens to be a string
-    assert(Array.isArray(codegenSummaryRequest.args.files), 'files is not an array');
+    assert(Array.isArray(codegenSummaryRequest.args.fileUpdates), 'fileUpdates is not an array');
     assert(Array.isArray(codegenSummaryRequest.args.contextPaths), 'contextPaths is not an array');
 
     if (codegenSummaryRequest.args.contextPaths.length > 0 && !disableContextOptimization) {
@@ -60,7 +65,7 @@ export async function promptService(generateContentFn, generateImageFn, codegenP
       // Monkey patch the initial getSourceCode, do not send parts of source code that are consider irrelevant
       getSourceCodeRequest.args = {
         filePaths: [
-          ...codegenSummaryRequest.args.files.map((file) => file.path),
+          ...codegenSummaryRequest.args.fileUpdates.map((file) => file.path),
           ...codegenSummaryRequest.args.contextPaths,
         ],
       };
@@ -77,7 +82,7 @@ export async function promptService(generateContentFn, generateImageFn, codegenP
 
     const result = [];
 
-    for (const file of codegenSummaryRequest.args.files) {
+    for (const file of codegenSummaryRequest.args.fileUpdates) {
       console.log('Collecting partial update for: ' + file.path + ' using tool: ' + file.updateToolName);
       console.log('- Prompt:', file.prompt);
       console.log('- Temperature', file.temperature);
@@ -93,7 +98,7 @@ export async function promptService(generateContentFn, generateImageFn, codegenP
         prompt.push({ type: 'user', text: file.prompt ?? messages.partialPromptTemplate(file.path) });
       }
 
-      if (vision) {
+      if (vision && file.contextImageAssets) {
         prompt.slice(-1)[0].images = file.contextImageAssets.map((path) => ({
           path,
           base64url: fs.readFileSync(path, 'base64'),
@@ -101,13 +106,14 @@ export async function promptService(generateContentFn, generateImageFn, codegenP
         }));
       }
 
-      let partialResult = await generateContentFn(
+      let partialRequest = [
         prompt,
         functionDefs,
         file.updateToolName,
         file.temperature ?? temperature,
         file.cheap === true,
-      );
+      ];
+      let partialResult = await generateContentFn(...partialRequest);
 
       let getSourceCodeCall = partialResult.find((call) => call.name === 'getSourceCode');
       assert(!getSourceCodeCall, 'Unexpected getSourceCode: ' + JSON.stringify(getSourceCodeCall));
@@ -141,6 +147,13 @@ export async function promptService(generateContentFn, generateImageFn, codegenP
             },
           });
         }
+      } else {
+        partialResult = await validateAndRecoverSingleResult(
+          partialRequest,
+          partialResult,
+          messages,
+          generateContentFn,
+        );
       }
 
       // Verify if patchFile is one of the functions called, and test if patch is valid and can be applied successfully
@@ -162,12 +175,14 @@ export async function promptService(generateContentFn, generateImageFn, codegenP
           console.log(`Patch could not be applied for ${filePath}. Retrying without patchFile function.`);
 
           // Rerun content generation without patchFile function
-          partialResult = await generateContentFn(
-            prompt,
-            functionDefs,
-            'updateFile',
-            file.temperature,
-            file.cheap === true,
+          partialRequest = [prompt, functionDefs, 'updateFile', file.temperature ?? temperature, file.cheap === true];
+          partialResult = await generateContentFn(...partialRequest);
+
+          partialResult = await validateAndRecoverSingleResult(
+            partialRequest,
+            partialResult,
+            messages,
+            generateContentFn,
           );
 
           let getSourceCodeCall = partialResult.find((call) => call.name === 'getSourceCode');
@@ -198,6 +213,56 @@ export async function promptService(generateContentFn, generateImageFn, codegenP
   }
 }
 
+async function validateAndRecoverSingleResult(
+  [prompt, functionDefs, requiredFunctionName, temperature, cheap],
+  result,
+  messages,
+  generateContentFn,
+) {
+  if (result.length !== 1) {
+    return result;
+  }
+
+  const call = result[0];
+  const validatorError = validateFunctionCall(call);
+  if (validatorError) {
+    console.log('Invalid function call', call, validatorError);
+
+    prompt.push(
+      { type: 'assistant', functionCalls: [call] },
+      {
+        type: 'user',
+        text: messages.invalidFunctionCall,
+        functionResponses: [
+          {
+            name: call.name,
+            call_id: call.id,
+            content: JSON.stringify({ args: call.args, error: validatorError }),
+            isError: true,
+          },
+        ],
+      },
+    );
+
+    console.log('Trying to recover...');
+    if (cheap) {
+      console.log('Disabling --cheap for recovery.');
+    }
+    result = await generateContentFn(prompt, functionDefs, requiredFunctionName, temperature, false);
+    console.log('Recover result:', result);
+
+    if (result.length === 1) {
+      const recoveryError = validateFunctionCall(result[0]);
+      assert(!recoveryError, 'Recovery failed');
+      console.log('Recovery was sucessful');
+    } else {
+      console.log('Unexpected number of function calls', result);
+    }
+  }
+
+  return result;
+}
+
 /**
  * Function to prepare messages for AI services
  */
@@ -216,5 +281,7 @@ function prepareMessages(prompt) {
     partialPromptTemplate(path) {
       return `Thank you for providing the summary, now suggest changes for the \`${path}\` file using appropriate tools.`;
     },
+    invalidFunctionCall:
+      'Function call was invalid, please analyze the error and respond with corrected function call.',
   };
 }

package/src/prompt/prompt-service.test.js

@@ -124,7 +124,7 @@ describe('promptService', () => {
       {
         name: 'codegenSummary',
         args: {
-          files: [{ path: 'test.js', updateToolName: 'patchFile' }],
+          fileUpdates: [{ path: 'test.js', updateToolName: 'patchFile' }],
           contextPaths: [],
           explanation: 'Mock summary',
         },
@@ -213,7 +213,7 @@ describe('promptService', () => {
       {
         name: 'codegenSummary',
         args: {
-          files: [
+          fileUpdates: [
             {
               path: 'test.js',
               updateToolName: 'updateFile',
@@ -296,7 +296,7 @@ describe('promptService', () => {
       {
         name: 'codegenSummary',
         args: {
-          files: [{ path: 'test.js', updateToolName: 'updateFile' }],
+          fileUpdates: [{ path: 'test.js', updateToolName: 'updateFile' }],
           contextPaths: ['context1.js', 'context2.js'],
           explanation: 'Mock summary with context',
         },
@@ -341,7 +341,7 @@ describe('promptService', () => {
       {
         name: 'codegenSummary',
         args: {
-          files: [{ path: 'test.js', updateToolName: 'updateFile' }],
+          fileUpdates: [{ path: 'test.js', updateToolName: 'updateFile' }],
           contextPaths: ['context1.js', 'context2.js'],
           explanation: 'Mock summary without context optimization',
         },
@@ -386,7 +386,7 @@ describe('promptService', () => {
       {
         name: 'codegenSummary',
         args: {
-          files: [{ path: 'test.js', updateToolName: 'updateFile' }],
+          fileUpdates: [{ path: 'test.js', updateToolName: 'updateFile' }],
           contextPaths: [],
           explanation: 'Mock summary with image generation',
         },
@@ -436,7 +436,7 @@ describe('promptService', () => {
       {
         name: 'codegenSummary',
         args: {
-          files: [{ path: 'test.js', updateToolName: 'updateFile' }],
+          fileUpdates: [{ path: 'test.js', updateToolName: 'updateFile' }],
           contextPaths: [],
           explanation: 'Mock summary with image generation failure',
         },
@@ -497,4 +497,105 @@ describe('promptService', () => {
     expect(vertexAi.generateContent).toHaveBeenCalledTimes(1);
     expect(result).toEqual(mockUnexpectedResponse);
   });
+
+  describe('validateAndRecoverSingleResult', () => {
+    it('should successfully recover from an invalid function call', async () => {
+      cliParams.vertexAi = true;
+      const mockInvalidCall = [
+        {
+          name: 'codegenSummary',
+          args: {
+            files: [{ path: 'test.js', updateToolName: 'patchFile' }],
+            contextPaths: [],
+            explanation: 'Mock summary',
+          },
+        },
+      ];
+      const mockValidCall = [
+        {
+          name: 'codegenSummary',
+          args: {
+            fileUpdates: [{ path: 'test.js', updateToolName: 'patchFile' }],
+            contextPaths: [],
+            explanation: 'Mock summary',
+          },
+        },
+      ];
+
+      vertexAi.generateContent
+        .mockResolvedValueOnce(mockInvalidCall)
+        .mockResolvedValueOnce(mockValidCall)
+        .mockResolvedValueOnce([
+          {
+            name: 'updateFile',
+            args: {
+              filePath: 'test.js',
+              newContent: 'console.log("Unexpected response");',
+            },
+          },
+        ]);
+
+      const result = await promptService(vertexAi.generateContent);
+
+      expect(vertexAi.generateContent).toHaveBeenCalledTimes(3);
+      expect(result).toEqual([
+        {
+          args: {
+            filePath: 'test.js',
+            newContent: 'console.log("Unexpected response");',
+          },
+          name: 'updateFile',
+        },
+      ]);
+    });
+
+    it('should handle unsuccessful recovery', async () => {
+      cliParams.vertexAi = true;
+      const mockInvalidCall = [
+        {
+          name: 'updateFile',
+          args: { filePath: 'test.js', invalidArg: 'This should not be here' },
+        },
+      ];
+
+      vertexAi.generateContent
+        .mockResolvedValueOnce([
+          {
+            name: 'codegenSummary',
+            args: {
+              fileUpdates: [{ path: 'test.js', updateToolName: 'patchFile' }],
+              contextPaths: [],
+              explanation: 'Mock summary',
+            },
+          },
+        ])
+        .mockResolvedValueOnce(mockInvalidCall)
+        .mockResolvedValueOnce(mockInvalidCall); // Second call also returns invalid result
+
+      await expect(promptService(vertexAi.generateContent)).rejects.toThrow('Recovery failed');
+
+      expect(vertexAi.generateContent).toHaveBeenCalledTimes(3);
+    });
+
+    it('should not attempt recovery for multiple valid function calls', async () => {
+      cliParams.vertexAi = true;
+      const mockValidCalls = [
+        {
+          name: 'updateFile',
+          args: { filePath: 'test1.js', newContent: 'console.log("File 1");' },
+        },
+        {
+          name: 'updateFile',
+          args: { filePath: 'test2.js', newContent: 'console.log("File 2");' },
+        },
+      ];
+
+      vertexAi.generateContent.mockResolvedValueOnce(mockValidCalls);
+
+      const result = await promptService(vertexAi.generateContent);
+
+      expect(vertexAi.generateContent).toHaveBeenCalledTimes(1);
+      expect(result).toEqual(mockValidCalls);
+    });
+  });
 });