add support for google gemini API

presubmit · Nov 10, 2024 · 57eeabc · 57eeabc
1 parent 97807d1
commit 57eeabc
Show file tree

Hide file tree

Showing 5 changed files with 125 additions and 99 deletions.
diff --git a/dist/index.js b/dist/index.js
diff --git a/package.json b/package.json
@@ -20,6 +20,7 @@
     "@actions/core": "^1.11.1",
     "@actions/github": "^6.0.0",
     "@ai-sdk/anthropic": "^0.0.56",
+    "@ai-sdk/google": "^0.0.55",
     "@ai-sdk/openai": "^0.0.72",
     "@octokit/action": "^7.0.0",
     "@octokit/plugin-retry": "^7.1.2",

diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/src/ai.ts b/src/ai.ts
@@ -1,6 +1,7 @@
 import { generateObject } from "ai";
 import { createAnthropic } from "@ai-sdk/anthropic";
 import { createOpenAI } from "@ai-sdk/openai";
+import { createGoogleGenerativeAI } from "@ai-sdk/google";
 import { z } from "zod";
 import { info } from "@actions/core";
 import config from "./config";
@@ -20,6 +21,15 @@ const LLM_MODELS = [
     name: "gpt-4o-mini",
     createAi: createOpenAI,
   },
+  // Google
+  {
+    name: "gemini-1.5-flash",
+    createAi: createGoogleGenerativeAI,
+  },
+  {
+    name: "gemini-1.5-flash-latest",
+    createAi: createGoogleGenerativeAI,
+  },
 ];
 
 export async function runPrompt({

diff --git a/src/pull_request.ts b/src/pull_request.ts
@@ -81,13 +81,13 @@ export async function handlePullRequest() {
   info(`generated pull request summary: ${summary.title}`);
 
   // Update PR title and description
-  await octokit.rest.pulls.update({
-    ...context.repo,
-    pull_number: pull_request.number,
-    title: summary.title,
-    body: summary.description,
-  });
-  info(`updated pull request title and description`);
+  //   await octokit.rest.pulls.update({
+  //     ...context.repo,
+  //     pull_number: pull_request.number,
+  //     title: summary.title,
+  //     body: summary.description,
+  //   });
+  //   info(`updated pull request title and description`);
 
   // Update overview comment with the walkthrough
   await octokit.rest.issues.updateComment({