Merge pull request #22 from dot-mario/main

Sitoi · web-flow · commit 7ba5bbd06d48 · 2025-02-11T16:59:40.000+08:00
✨ feat(config): add gemini support as ai provider
diff --git a/README.md b/README.md
@@ -27,7 +27,7 @@ Use Azure/OpenAI API to review Git changes, generate conventional commit message
 
 ## ✨ Features
 
-- 🤯 Support generating commit messages based on git diffs using ChatGPT / Azure API.
+- 🤯 Support generating commit messages based on git diffs using ChatGPT / Azure API and Gemini API.
 - 🗺️ Support multi-language commit messages.
 - 😜 Support adding Gitmoji.
 - 🛠️ Support custom system prompt.
@@ -59,15 +59,19 @@ Use Azure/OpenAI API to review Git changes, generate conventional commit message
 
 In the VSCode settings, locate the "ai-commit" configuration options and configure them as needed:
 
-| Configuration      |  Type  | Default | Required |                                                  Notes                                                  |
-| :----------------- | :----: | :-----: | :------: | :-----------------------------------------------------------------------------------------------------: |
-| OPENAI_API_KEY     | string |  None   |   Yes    |                      [OpenAI token](https://platform.openai.com/account/api-keys)                       |
-| OPENAI_BASE_URL    | string |  None   |    No    |           If using Azure, use: https://{resource}.openai.azure.com/openai/deployments/{model}           |
-| OPENAI_MODEL       | string | gpt-4o  |   Yes    | OpenAI MODEL,you can select a model from the list by running the `Show Available OpenAI Models` command |
-| AZURE_API_VERSION  | string |  None   |    No    |                                            AZURE_API_VERSION                                            |
-| AI_COMMIT_LANGUAGE | string |   en    |   Yes    |                                          Supports 19 languages                                          |
-| SYSTEM_PROMPT      | string |  None   |    No    |                                          Custom system prompt                                           |
-| OPENAI_TEMPERATURE | number |   0.7   |    No    | Controls randomness in the output. Range: 0-2. Lower values: more focused, Higher values: more creative |
+| Configuration      |  Type  |       Default        | Required |                                                       Notes                                                        |
+| :----------------- | :----: | :------------------: | :------: | :----------------------------------------------------------------------------------------------------------------: |
+| AI_PROVIDER        | string |        openai        |   Yes    |                                     Select AI Provider: `openai` or `gemini`.                                      |
+| OPENAI_API_KEY     | string |         None         |   Yes    |    Required when `AI Provider` is set to `OpenAI`. [OpenAI token](https://platform.openai.com/account/api-keys)    |
+| OPENAI_BASE_URL    | string |         None         |    No    |                If using Azure, use: https://{resource}.openai.azure.com/openai/deployments/{model}                 |
+| OPENAI_MODEL       | string |        gpt-4o        |   Yes    |      OpenAI MODEL, you can select a model from the list by running the `Show Available OpenAI Models` command      |
+| AZURE_API_VERSION  | string |         None         |    No    |                                                 AZURE_API_VERSION                                                  |
+| OPENAI_TEMPERATURE | number |         0.7          |    No    |      Controls randomness in the output. Range: 0-2. Lower values: more focused, Higher values: more creative       |
+| GEMINI_API_KEY     | string |         None         |   Yes    |     Required when `AI Provider` is set to `Gemini`. [Gemini API key](https://makersuite.google.com/app/apikey)     |
+| GEMINI_MODEL       | string | gemini-2.0-flash-001 |   Yes    |                       Gemini MODEL.  Currently, model selection is limited to configuration.                       |
+| GEMINI_TEMPERATURE | number |         0.7          |    No    | Controls randomness in the output. Range: 0-2 for Gemini. Lower values: more focused, Higher values: more creative |
+| AI_COMMIT_LANGUAGE | string |          en          |   Yes    |                                               Supports 19 languages                                                |
+| SYSTEM_PROMPT      | string |         None         |    No    |                                               Custom system prompt                                                 |
 
 ## ⌨️ Local Development
 
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -124,6 +124,32 @@
           "minimum": 0,
           "maximum": 2,
           "description": "OpenAI temperature setting (0-2). Higher values make output more random, lower values more deterministic."
+        },
+        "ai-commit.AI_PROVIDER": {
+          "type": "string",
+          "default": "openai",
+          "description": "AI Provider to use (OpenAI or Gemini)",
+          "enum": [
+            "openai",
+            "gemini"
+          ]
+        },
+        "ai-commit.GEMINI_API_KEY": {
+          "type": "string",
+          "default": "",
+          "description": "Gemini API Key"
+        },
+        "ai-commit.GEMINI_MODEL": {
+          "type": "string",
+          "default": "gemini-2.0-flash-001",
+          "description": "Gemini Model to use"
+        },
+        "ai-commit.GEMINI_TEMPERATURE": {
+          "type": "number",
+          "default": 0.7,
+          "minimum": 0,
+          "maximum": 2,
+          "description": "Gemini temperature setting (0-2). Controls randomness."
         }
       },
       "title": "AI Commit"
@@ -143,9 +169,9 @@
     "compile": "webpack",
     "compile-tests": "tsc -p . --outDir out",
     "lint": "eslint src --ext ts",
-    "package": "npm vsce package --no-dependencies",
+    "package": "vsce package --no-dependencies",
     "pretest": "npm run compile-tests && npm run compile && npm run lint",
-    "publish": "npm vsce publish --no-dependencies",
+    "publish": "vsce publish --no-dependencies",
     "test": "node ./out/test/runTest.js",
     "vscode:prepublish": "npm run build",
     "watch": "webpack --watch",
@@ -169,9 +195,10 @@
     "webpack-cli": "^5.0.1"
   },
   "dependencies": {
+    "@google/generative-ai": "^0.21.0",
+    "fs-extra": "^11.0.4",
     "openai": "^4.14.2",
-    "simple-git": "^3.17.0",
-    "fs-extra": "^11.0.4"
+    "simple-git": "^3.17.0"
   },
   "resolutions": {
     "@types/node": "16.x"
diff --git a/src/commands.ts b/src/commands.ts
@@ -19,7 +19,7 @@ export class CommandManager {
     // Show available OpenAI models
     this.registerCommand('ai-commit.showAvailableModels', async () => {
       const configManager = ConfigurationManager.getInstance();
-      const models = await configManager.getAvailableModels();
+      const models = await configManager.getAvailableOpenAIModels();
       const selected = await vscode.window.showQuickPick(models, {
         placeHolder: 'Please select a model'
       });
@@ -29,6 +29,27 @@ export class CommandManager {
         await config.update('OPENAI_MODEL', selected, vscode.ConfigurationTarget.Global);
       }
     });
+
+    /**
+     * @deprecated
+     * This function is deprecated because Gemini API does not currently support listing models via API.
+     * 
+     * Show available Gemini models
+     */
+    /*
+    this.registerCommand('ai-commit.showAvailableGeminiModels', async () => {
+      const configManager = ConfigurationManager.getInstance();
+      const models = await configManager.getAvailableGeminiModels(); // Use the updated function
+      const selected = await vscode.window.showQuickPick(models, {
+        placeHolder: 'Please select a Gemini model'
+      });
+
+      if (selected) {
+        const config = vscode.workspace.getConfiguration('ai-commit');
+        await config.update('GEMINI_MODEL', selected, vscode.ConfigurationTarget.Global);
+      }
+    });
+    */
   }
 
   private registerCommand(command: string, handler: (...args: any[]) => any) {
diff --git a/src/config.ts b/src/config.ts
@@ -1,5 +1,6 @@
 import * as vscode from 'vscode';
 import { createOpenAIApi } from './openai-utils';
+import { createGeminiAPIClient } from './gemini-utils';
 
 /**
  * Configuration keys used in the AI commit extension.
@@ -19,7 +20,12 @@ export enum ConfigKeys {
   AZURE_API_VERSION = 'AZURE_API_VERSION',
   AI_COMMIT_LANGUAGE = 'AI_COMMIT_LANGUAGE',
   SYSTEM_PROMPT = 'AI_COMMIT_SYSTEM_PROMPT',
-  OPENAI_TEMPERATURE = 'OPENAI_TEMPERATURE'
+  OPENAI_TEMPERATURE = 'OPENAI_TEMPERATURE',
+  
+  GEMINI_API_KEY = 'GEMINI_API_KEY',
+  GEMINI_MODEL = 'GEMINI_MODEL',
+  GEMINI_TEMPERATURE = 'GEMINI_TEMPERATURE',
+  AI_PROVIDER = 'AI_PROVIDER',
 }
 
 /**
@@ -39,7 +45,7 @@ export class ConfigurationManager {
 
         if (event.affectsConfiguration('ai-commit.OPENAI_BASE_URL') ||
           event.affectsConfiguration('ai-commit.OPENAI_API_KEY')) {
-          this.updateModelList();
+          this.updateOpenAIModelList();
         }
       }
     });
@@ -67,13 +73,13 @@ export class ConfigurationManager {
   /**
    * Updates the list of available OpenAI models.
    */
-  private async updateModelList() {
+  private async updateOpenAIModelList() {
     try {
       const openai = createOpenAIApi();
       const models = await openai.models.list();
 
       // Save available models to extension state
-      await this.context.globalState.update('availableModels', models.data.map(model => model.id));
+      await this.context.globalState.update('availableOpenAIModels', models.data.map(model => model.id));
 
       // Get the current selected model
       const config = vscode.workspace.getConfiguration('ai-commit');
@@ -91,12 +97,60 @@ export class ConfigurationManager {
 
   /**
    * Retrieves the list of available OpenAI models.
-   * @returns {Promise<string[]>} The list of available models.
+   * @returns {Promise<string[]>} The list of available OpenAI models.
    */
-  public async getAvailableModels(): Promise<string[]> {
-    if (!this.context.globalState.get<string[]>('availableModels')) {
-      await this.updateModelList();
+  public async getAvailableOpenAIModels(): Promise<string[]> {
+    if (!this.context.globalState.get<string[]>('availableOpenAIModels')) {
+      await this.updateOpenAIModelList();
     }
-    return this.context.globalState.get<string[]>('availableModels', []);
+    return this.context.globalState.get<string[]>('availableOpenAIModels', []);
   }
+
+  /**
+   * @deprecated
+   * This function is deprecated because Gemini API does not currently support listing models via API.
+   * We have to wait for this feature to be updated to the gemini library at some point, or find another way.
+   * 
+   * Updates the list of available Gemini models.
+   */
+  /*
+  private async updateGeminiModelList() {
+    try {
+      const geminiAPI = createGeminiAPIClient();
+      const modelListResponse = await geminiAPI.listModels(); // Gemini API does not currently have a function to get a list of models
+      const availableModels = modelListResponse.models.map(model => model.name);
+
+      // Save available Gemini models to extension global state
+      await this.context.globalState.update('availableGeminiModels', availableModels);
+
+      // Get the currently selected Gemini model
+      const config = vscode.workspace.getConfiguration('ai-commit');
+      const currentModel = config.get<string>('GEMINI_MODEL');
+
+      // If the current selected Gemini model is not in the available list, set it to a default value
+      if (currentModel && !availableModels.includes(currentModel)) {
+        await config.update('GEMINI_MODEL', 'gemini-2.0-flash-001', vscode.ConfigurationTarget.Global);
+      }
+
+    } catch (error) {
+      console.error('Failed to fetch Gemini models:', error);
+    }
+  }
+  */
+
+  /**
+   * @deprecated
+   * This function is deprecated because Gemini API does not currently support listing models via API.
+   * 
+   * Retrieves the list of available Gemini models.
+   * @returns {Promise<string[]>} The list of available Gemini models.
+   */
+  /*
+  public async getAvailableGeminiModels(): Promise<string[]> {
+    if (!this.context.globalState.get<string[]>('availableGeminiModels')) {
+      await this.updateGeminiModelList();
+    }
+    return this.context.globalState.get<string[]>('availableGeminiModels', []);
+  }
+  */
 }
diff --git a/src/gemini-utils.ts b/src/gemini-utils.ts
@@ -0,0 +1,64 @@
+import { GoogleGenerativeAI } from "@google/generative-ai";
+import { ConfigKeys, ConfigurationManager } from './config';
+
+/**
+ * Creates and returns a Gemini API configuration object.
+ * @returns {Object} - The Gemini API configuration object.
+ * @throws {Error} - Throws an error if the API key is missing or empty.
+ */
+function getGeminiConfig() {
+  const configManager = ConfigurationManager.getInstance();
+  const apiKey = configManager.getConfig<string>(ConfigKeys.GEMINI_API_KEY);
+
+  if (!apiKey) {
+    throw new Error('The GEMINI_API_KEY environment variable is missing or empty.');
+  }
+
+  const config: {
+    apiKey: string;
+  } = {
+    apiKey
+  };
+
+  return config;
+}
+
+/**
+ * Creates and returns a Gemini API instance.
+ * @returns {GoogleGenerativeAI} - The Gemini API instance.
+ */
+export function createGeminiAPIClient() {
+  const config = getGeminiConfig();
+  return new GoogleGenerativeAI(config.apiKey);
+}
+
+/**
+ * Sends a chat completion request to the Gemini API.
+ * @param {any[]} messages - The messages to send to the API.
+ * @returns {Promise<string>} - A promise that resolves to the API response.
+ */
+export async function GeminiAPI(messages: any[]) {
+  try {
+    const gemini = createGeminiAPIClient();
+    const configManager = ConfigurationManager.getInstance();
+    const modelName = configManager.getConfig<string>(ConfigKeys.GEMINI_MODEL);
+    const temperature = configManager.getConfig<number>(ConfigKeys.GEMINI_TEMPERATURE, 0.7);
+
+    const model = gemini.getGenerativeModel({ model: modelName });
+    const chat = model.startChat({
+      generationConfig: {
+        temperature: temperature,
+      },
+    });
+
+    const result = await chat.sendMessage(messages.map(msg => msg.content));
+    const response = result.response;
+    const text = response.text();
+
+    return text;
+
+  } catch (error) {
+    console.error('Gemini API call failed:', error);
+    throw error;
+  }
+}
diff --git a/src/generate-commit-msg.ts b/src/generate-commit-msg.ts
