feat: add Gemini LLM provider

DiogoRibeiro7 · DiogoRibeiro7 · commit dbabd4d617db · 2025-06-11T22:34:11.000+01:00
diff --git a/README.md b/README.md
@@ -7,7 +7,7 @@ A GitHub Action that scans your codebase for inline TODOs, FIXMEs, and BUG comme
 ## 🚀 Features
 
 - ✅ Detects `TODO`, `FIXME`, `BUG`, and `HACK` comments
-- ✅ Supports multiple languages: `.ts`, `.js`, `.py`, `.go`, `.html`, etc.
+- ✅ Supports many languages: `.ts`, `.js`, `.py`, `.go`, `.c`, `.cpp`, `.rs`, `.html`, `.yaml`, etc.
 - ✅ Skips common directories like `node_modules`, `dist`, and `coverage`
 - ✅ Extracts metadata like `priority`, `due`, etc.
 - ✅ Parses structured tags (`@assignee`, `#module`, `key=value`)
@@ -17,6 +17,8 @@ A GitHub Action that scans your codebase for inline TODOs, FIXMEs, and BUG comme
 - ✅ Supports custom label colors and descriptions via JSON config
 - ✅ Custom templates for issue titles and bodies
 - ✅ LLM-powered issue title and body generation
+- ✅ Automatic retry logic for OpenAI API calls
+- ✅ Supports multiple LLM providers: OpenAI or Gemini
 - ✅ Command-line interface for local usage
 - ✅ Optional Jira synchronization
 
@@ -55,8 +57,12 @@ jobs:
         with:
           repo-token: ${{ secrets.GITHUB_TOKEN }}
           limit: 5
+          llm: true
+          llm-provider: openai # or 'gemini'
 ```
 
+Set `OPENAI_API_KEY` or `GEMINI_API_KEY` secrets based on your chosen provider.
+
 ### 2. Run the CLI locally
 
 Use the bundled command-line interface to scan a directory on your machine and
diff --git a/action.yml b/action.yml
@@ -39,6 +39,11 @@ inputs:
     description: Use LLM to generate issue titles and bodies
     default: 'false'
 
+  llm-provider:
+    required: false
+    description: LLM provider to use (`openai` or `gemini`)
+    default: openai
+
   openai-api-key:
     required: false
     description: 'OpenAI API key used when `llm` is true'
@@ -48,6 +53,15 @@ inputs:
     description: OpenAI model to use (e.g., `gpt-3.5-turbo`, `gpt-4`)
     default: gpt-3.5-turbo
 
+  gemini-api-key:
+    required: false
+    description: Gemini API key used when `llm-provider` is `gemini`
+
+  gemini-model:
+    required: false
+    description: Gemini model to use (e.g., `gemini-1.5-pro`)
+    default: gemini-1.5-pro
+
   sync-to-jira:
     required: false
     default: 'false'
diff --git a/package.json b/package.json
@@ -26,6 +26,7 @@
   "dependencies": {
     "@actions/core": "^1.10.0",
     "@actions/github": "^6.0.1",
+    "@google/genai": "^1.4.0",
     "@jest/globals": "^29.7.0",
     "@octokit/rest": "^21.1.1",
     "@types/jest": "^29.5.14",
diff --git a/src/ActionMain.ts b/src/ActionMain.ts
@@ -23,11 +23,20 @@ async function run(): Promise<void> {
     const workspace = process.env.GITHUB_WORKSPACE || '.';
 
     // LLM support
-    process.env.OPENAI_API_KEY = core.getInput('openai-api-key') || process.env.OPENAI_API_KEY;
+    const llmProvider = core.getInput('llm-provider') || 'openai';
+    process.env.LLM_PROVIDER = llmProvider;
+    if (llmProvider === 'gemini') {
+      process.env.GEMINI_API_KEY = core.getInput('gemini-api-key') || process.env.GEMINI_API_KEY;
+    } else {
+      process.env.OPENAI_API_KEY = core.getInput('openai-api-key') || process.env.OPENAI_API_KEY;
+    }
     const useLLM = core.getInput('llm') === 'true';
-    if (useLLM && !process.env.OPENAI_API_KEY) {
+    if (useLLM && llmProvider === 'openai' && !process.env.OPENAI_API_KEY) {
       core.warning('⚠️ LLM is enabled, but OPENAI_API_KEY is not set.');
     }
+    if (useLLM && llmProvider === 'gemini' && !process.env.GEMINI_API_KEY) {
+      core.warning('⚠️ LLM is enabled, but GEMINI_API_KEY is not set.');
+    }
 
     const useStructured = core.getInput('structured') === 'true';
 
diff --git a/src/core/llm/generateIssueContent.ts b/src/core/llm/generateIssueContent.ts
@@ -1,13 +1,13 @@
 // src/core/llm/generateIssueContent.ts
 import { TodoItem } from '../../parser/types';
-import OpenAI from 'openai';
 import * as core from '@actions/core';
+import { chatCompletionWithRetry } from './llmClient';
 
-const openai = new OpenAI({
-    apiKey: core.getInput('openai-api-key'), // correto agora
-  });
-
-const model = core.getInput('openai-model') || 'gpt-3.5-turbo';
+const provider = core.getInput('llm-provider') || 'openai';
+const model =
+  provider === 'gemini'
+    ? core.getInput('gemini-model') || 'gemini-1.5-pro'
+    : core.getInput('openai-model') || 'gpt-3.5-turbo';
 
 export async function generateIssueTitleAndBodyLLM(todo: TodoItem): Promise<{ title: string; body: string }> {
   const prompt = `
@@ -27,17 +27,20 @@ TITLE: <title>
 BODY:
 <detailed body>
 `;
-  // 👇 Adiciona aqui
-  core.debug(`[DEBUG] OpenAI key starts with: ${process.env.OPENAI_API_KEY?.slice(0, 5)}`);
+  core.debug(`[DEBUG] LLM provider: ${provider}`);
+  if (provider === 'openai') {
+    core.debug(`[DEBUG] OpenAI key starts with: ${process.env.OPENAI_API_KEY?.slice(0, 5)}`);
+  } else {
+    core.debug(`[DEBUG] Gemini key starts with: ${process.env.GEMINI_API_KEY?.slice(0, 5)}`);
+  }
   core.debug(`[DEBUG] Using model: ${model}`);
-  core.debug('[DEBUG] Sending prompt to OpenAI...');
+  core.debug('[DEBUG] Sending prompt to LLM...');
 try {
-  const response = await openai.chat.completions.create({
+  const response = await chatCompletionWithRetry({
     model,
     messages: [{ role: 'user', content: prompt }],
     temperature: 0.4,
   });
-  // TODO(priority=high): improve retry logic for API errors
   const result = response.choices[0].message?.content || '';
   const match = result.match(/TITLE:\s*(.+?)\s*BODY:\s*([\s\S]*)/i);
 
@@ -48,7 +51,7 @@ try {
   const [, title, body] = match;
   return { title: title.trim(), body: body.trim() };
 } catch (err: any) {
-  console.error('[ERROR] OpenAI call failed:', err);
+  console.error('[ERROR] LLM call failed:', err);
   throw err;
 }
 }
diff --git a/src/core/llm/llmClient.ts b/src/core/llm/llmClient.ts
@@ -0,0 +1,54 @@
+// src/core/llm/openaiClient.ts
+import OpenAI from 'openai';
+import { GoogleGenAI } from '@google/genai';
+import * as core from '@actions/core';
+
+const provider = core.getInput('llm-provider') || 'openai';
+
+export const openai = new OpenAI({
+  apiKey: core.getInput('openai-api-key'),
+});
+
+export const gemini = new GoogleGenAI({
+  apiKey: core.getInput('gemini-api-key'),
+});
+
+/**
+ * Wraps `openai.chat.completions.create` with simple retry logic.
+ * Retries on failure with exponential backoff.
+ *
+ * @param params Parameters forwarded to OpenAI
+ * @param maxRetries Maximum number of attempts before throwing
+ */
+export interface ChatCompletionParams {
+  model: string;
+  messages: { role: string; content: string }[];
+  temperature?: number;
+}
+
+export async function chatCompletionWithRetry(
+  params: ChatCompletionParams,
+  maxRetries = 3
+): Promise<{ choices: { message: { content: string } }[] }> {
+  let attempt = 0;
+  for (;;) {
+    try {
+      if (provider === 'gemini') {
+        const prompt = params.messages.map(m => m.content).join('\n');
+        const response = await gemini.models.generateContent({
+          model: params.model,
+          contents: prompt,
+          generationConfig: { temperature: params.temperature },
+        } as any);
+        return { choices: [{ message: { content: (response as any).text } }] } as any;
+      }
+      return (await openai.chat.completions.create(params as any)) as any;
+    } catch (err) {
+      attempt++;
+      if (attempt > maxRetries) throw err;
+      const delay = Math.min(1000 * 2 ** attempt, 5000);
+      core.warning(`LLM request failed (attempt ${attempt}). Retrying in ${delay}ms...`);
+      await new Promise(res => setTimeout(res, delay));
+    }
+  }
+}
diff --git a/src/core/llm/openaiClient.ts b/src/core/llm/openaiClient.ts
diff --git a/src/parser/extractTodos.ts b/src/parser/extractTodos.ts
@@ -4,8 +4,8 @@ import { TodoItem } from './types';
 import { normalizeTag } from '../utils/isTextFile';
 
 const COMMENT_PATTERNS = [
-  { ext: ['.ts', '.js', '.java', '.go'], pattern: /^\s*\/\/\s*(.*)$/ },
-  { ext: ['.py', '.sh', '.rb'], pattern: /^\s*#\s*(.*)$/ },
+  { ext: ['.ts', '.js', '.java', '.go', '.c', '.cpp', '.cs', '.rs', '.php', '.h', '.hpp'], pattern: /^\s*\/\/\s*(.*)$/ },
+  { ext: ['.py', '.sh', '.rb', '.yaml', '.yml'], pattern: /^\s*#\s*(.*)$/ },
   { ext: ['.html', '.xml'], pattern: /<!--\s*(.*?)\s*-->/ }
 ];
 
diff --git a/src/parser/extractTodosFromContent.ts b/src/parser/extractTodosFromContent.ts
@@ -4,8 +4,8 @@ import { TodoItem } from './types';
 import { normalizeTag } from '../utils/isTextFile';
 
 const COMMENT_PATTERNS = [
-  { ext: ['.ts', '.js', '.java', '.go'], pattern: /^\s*\/\/\s*(.*)$/ },
-  { ext: ['.py', '.sh', '.rb'], pattern: /^\s*#\s*(.*)$/ },
+  { ext: ['.ts', '.js', '.java', '.go', '.c', '.cpp', '.cs', '.rs', '.php', '.h', '.hpp'], pattern: /^\s*\/\/\s*(.*)$/ },
+  { ext: ['.py', '.sh', '.rb', '.yaml', '.yml'], pattern: /^\s*#\s*(.*)$/ },
   { ext: ['.html', '.xml'], pattern: /<!--\s*(.*?)\s*-->/ }
 ];
 
diff --git a/src/parser/extractTodosFromDir.ts b/src/parser/extractTodosFromDir.ts
@@ -3,7 +3,7 @@ import path from 'path';
 import { extractTodosFromFile } from './extractTodos';
 import { TodoItem } from './types';
 
-const SUPPORTED_EXTENSIONS = ['.ts', '.js', '.py', '.go', '.java', '.rb', '.sh', '.html', '.xml'];
+const SUPPORTED_EXTENSIONS = ['.ts', '.js', '.py', '.go', '.java', '.rb', '.sh', '.c', '.cpp', '.cs', '.rs', '.php', '.h', '.hpp', '.html', '.xml', '.yaml', '.yml'];
 const IGNORED_DIRS = ['node_modules', 'dist', 'coverage'];
 
 export function extractTodosFromDir(dirPath: string): TodoItem[] {
diff --git a/src/utils/isTextFile.ts b/src/utils/isTextFile.ts
@@ -5,8 +5,8 @@
  * Useful for filtering files before parsing for TODOs.
  */
 export function isTextFile(filename: string): boolean {
-    return /\.(ts|js|jsx|tsx|py|rb|sh|go|java|html|css|json|md|txt|xml|yaml|yml)$/i.test(filename);
-  }
+  return /\.(ts|js|jsx|tsx|py|rb|sh|go|java|c|cpp|cs|rs|php|h|hpp|html|css|json|md|txt|xml|yaml|yml)$/i.test(filename);
+}
   
   /**
    * List of known multilingual aliases for TODO-related tags.
diff --git a/tests/commentPatterns.test.ts b/tests/commentPatterns.test.ts
@@ -2,9 +2,9 @@ import { describe, it, expect } from 'vitest';
 import { extractTodosFromString } from '../src/parser/extractTodosFromContent';
 
 describe('extractTodosFromString - comment support by extension', () => {
-  it('extracts from JS-style (//) for .js/.ts/.go/.java', () => {
+  it('extracts from JS-style (//) for many languages', () => {
     const code = `// TODO: js comment\n// BUG: broken`;
-    const extensions = ['.js', '.ts', '.go', '.java'];
+    const extensions = ['.js', '.ts', '.go', '.java', '.c', '.cpp', '.rs'];
 
     for (const ext of extensions) {
       const todos = extractTodosFromString(code, ext);
@@ -14,9 +14,9 @@ describe('extractTodosFromString - comment support by extension', () => {
     }
   });
 
-  it('extracts from Python-style (#) for .py/.sh/.rb', () => {
+  it('extracts from Python-style (#) for .py/.sh/.rb/.yaml', () => {
     const code = `# TODO: python comment\n# FIXME: fix me`;
-    const extensions = ['.py', '.sh', '.rb'];
+    const extensions = ['.py', '.sh', '.rb', '.yaml'];
 
     for (const ext of extensions) {
       const todos = extractTodosFromString(code, ext);
diff --git a/tests/generateIssueContent.test.ts b/tests/generateIssueContent.test.ts
@@ -1,36 +1,37 @@
 // tests/generateIssueContent.test.ts
 import { describe, it, expect, vi, beforeEach } from 'vitest';
-import * as core from '@actions/core';
 import OpenAI from 'openai';
-import { generateIssueTitleAndBodyLLM } from '../src/core/llm/generateIssueContent';
+
+async function loadGenerator(provider: string) {
+  vi.resetModules();
+  const coreModule = await import('@actions/core');
+  (coreModule.getInput as any) = vi.fn((key: string) => {
+    if (key === 'llm-provider') return provider;
+    if (key === 'openai-api-key') return 'fake-key';
+    if (key === 'openai-model') return 'gpt-3.5-turbo';
+    if (key === 'gemini-model') return 'gemini-1.5-pro';
+    return '';
+  });
+  return await import('../src/core/llm/generateIssueContent');
+}
 import { TodoItem } from '../src/parser/types';
 
 vi.mock('@actions/core');
 vi.mock('openai');
+vi.mock('@google/genai');
 
 describe('generateIssueTitleAndBodyLLM', () => {
-  const mockCore = core as unknown as { getInput: (key: string) => string };
   const mockCreate = vi.fn();
 
   beforeEach(() => {
     vi.clearAllMocks();
-
-    // Mock `core.getInput`
-    (mockCore.getInput as any) = vi.fn((key: string) => {
-      if (key === 'openai-api-key') return 'fake-key';
-      if (key === 'openai-model') return 'gpt-3.5-turbo';
-      return '';
-    });
-
-    // Mock `openai.chat.completions.create`
     (OpenAI as unknown as any).prototype.chat = {
-      completions: {
-        create: mockCreate,
-      },
+      completions: { create: mockCreate },
     };
   });
 
   it('should return title and body when OpenAI responds correctly', async () => {
+    const { generateIssueTitleAndBodyLLM } = await loadGenerator('openai');
     const todo: TodoItem = {
       tag: 'TODO',
       text: 'Refactor this function',
@@ -58,6 +59,7 @@ describe('generateIssueTitleAndBodyLLM', () => {
   });
 
   it('should throw if response format is invalid', async () => {
+    const { generateIssueTitleAndBodyLLM } = await loadGenerator('openai');
     const todo: TodoItem = {
       tag: 'TODO',
       text: 'Missing format',
@@ -78,6 +80,9 @@ describe('generateIssueTitleAndBodyLLM', () => {
   });
 
   it('should throw if OpenAI call fails', async () => {
+    const { generateIssueTitleAndBodyLLM } = await loadGenerator('openai');
+    vi.useFakeTimers();
+
     const todo: TodoItem = {
       tag: 'TODO',
       text: 'Handle error',
@@ -88,7 +93,29 @@ describe('generateIssueTitleAndBodyLLM', () => {
 
     mockCreate.mockRejectedValue(new Error('API failure'));
 
-    await expect(() => generateIssueTitleAndBodyLLM(todo)).rejects.toThrow('API failure');
+    const promise = generateIssueTitleAndBodyLLM(todo);
+    const expectation = expect(promise).rejects.toThrow('API failure');
+    await vi.runAllTimersAsync();
+    await expectation;
+    vi.useRealTimers();
+  });
+
+  it('should work with Gemini provider', async () => {
+    const { generateIssueTitleAndBodyLLM } = await loadGenerator('gemini');
+    const todo: TodoItem = {
+      tag: 'TODO',
+      text: 'Gemini task',
+      file: 'src/file.ts',
+      line: 5,
+      metadata: {},
+    };
+
+    (await import('@google/genai')).GoogleGenAI.prototype.models = {
+      generateContent: vi.fn().mockResolvedValue({ text: 'TITLE: G\nBODY: B' }),
+    } as any;
+
+    const result = await generateIssueTitleAndBodyLLM(todo);
+    expect(result).toEqual({ title: 'G', body: 'B' });
   });
 });
 
diff --git a/tests/isTextFile.test.ts b/tests/isTextFile.test.ts
@@ -0,0 +1,15 @@
+import { describe, it, expect } from 'vitest';
+import { isTextFile } from '../src/utils/isTextFile';
+
+describe('isTextFile', () => {
+  it('returns true for supported extensions', () => {
+    const files = ['main.c', 'module.cpp', 'program.rs', 'config.yaml'];
+    for (const f of files) {
+      expect(isTextFile(f)).toBe(true);
+    }
+  });
+
+  it('returns false for unsupported extensions', () => {
+    expect(isTextFile('image.png')).toBe(false);
+  });
+});
diff --git a/tests/llmClient.test.ts b/tests/llmClient.test.ts
diff --git a/yarn.lock b/yarn.lock
