Test

.github/workflows/test-action.yml

@@ -53,15 +53,10 @@ jobs:
           github_repository: ${{ github.repository }}
           github_pull_request_number: ${{ github.event.pull_request.number }}
           git_commit_hash: ${{ github.event.pull_request.head.sha }}
-          model: "gemini-1.5-pro-latest"
-          temperature: "0.1"
-          max_tokens: "512"
-          top_p: "1"
-          frequency_penalty: "0.0"
-          presence_penalty: "0.0"
+          model: "gemini-1.0-pro-latest"
           pull_request_diff: |-
             ${{ steps.get_diff.outputs.pull_request_diff }}
           pull_request_chunk_size: "3500"
           extra_prompt: |-
-            You are very familiar with python too.
+            You are very familiar with python too. Answer in Brazilian Portuguese!
           log_level: "DEBUG"

action.yml

@@ -42,26 +42,6 @@ inputs:
     description: "Extra prompt for GPT"
     required: false
     default: ""
-  temperature:
-    description: "Temperature for GPT"
-    required: false
-    default: "0.7"
-  max_tokens:
-    description: "Max tokens for GPT"
-    required: false
-    default: "256"
-  top_p:
-    description: "Top p for GPT"
-    required: false
-    default: "1"
-  frequency_penalty:
-    description: "Frequency penalty for GPT"
-    required: false
-    default: "0.0"
-  presence_penalty:
-    description: "Presence penalty for GPT"
-    required: false
-    default: "0.0"
   pull_request_diff:
     description: "Pull request diff"
     required: true
@@ -98,10 +78,5 @@ runs:
   args:
     - "--model=${{ inputs.model }}"
     - "--extra-prompt=${{ inputs.extra_prompt }}"
-    - "--temperature=${{ inputs.temperature }}"
-    - "--max-tokens=${{ inputs.max_tokens }}"
-    - "--top-p=${{ inputs.top_p }}"
-    - "--frequency-penalty=${{ inputs.frequency_penalty }}"
-    - "--presence-penalty=${{ inputs.presence_penalty }}"
     - "--diff-chunk-size=${{ inputs.pull_request_chunk_size }}"
     - "--diff=${{ inputs.pull_request_diff }}"

entrypoint.py

@@ -14,7 +14,7 @@
 from typing import List
 
 import click
-import openai
+import google.generativeai as genai
 import requests
 from loguru import logger
 
@@ -50,7 +50,7 @@ def get_summarize_prompt() -> str:
     """Get a prompt template"""
     template = """
     This is a pull request of a set of reviews of a pull request.
-    Those are generated by OpenAI's GPT.
+    Those are generated by Gemini AI.
     Can you summarized them?
     It would be good to focus on highlighting issues and providing suggestions to improve the pull request.
     """
@@ -99,44 +99,34 @@ def get_review(
 ):
     """Get a review"""
     # Chunk the prompt
+    genai_model = genai.GenerativeModel(model)
     review_prompt = get_review_prompt(extra_prompt=extra_prompt)
     chunked_diff_list = chunk_string(input_string=diff, chunk_size=prompt_chunk_size)
     # Get summary by chunk
     chunked_reviews = []
     for chunked_diff in chunked_diff_list:
-        response = openai.ChatCompletion.create(
-            model=model,
-            messages=[
-                {"role": "system", "content": review_prompt},
-                {"role": "user", "content": chunked_diff},
-            ],
-            temperature=temperature,
-            max_tokens=max_tokens,
-            top_p=top_p,
-            frequency_penalty=frequency_penalty,
-            presence_penalty=presence_penalty
-        )
-        review_result = response.choices[0].message["content"]
+        prompt = str(f"""
+        {str(review_prompt)}
+
+        ```{str(chunked_diff)}```""")
+
+        response = genai_model.generate_content(prompt)
+        review_result = response.text
         chunked_reviews.append(review_result)
     # If the chunked reviews are only one, return it
     if len(chunked_reviews) == 1:
         return chunked_reviews, chunked_reviews[0]
 
     # Summarize the chunked reviews
     summarize_prompt = get_summarize_prompt()
-    response = openai.ChatCompletion.create(
-        model=model,
-        messages=[
-            {"role": "system", "content": summarize_prompt},
-            {"role": "user", "content": "\n".join(chunked_reviews)},
-        ],
-        temperature=temperature,
-        max_tokens=max_tokens,
-        top_p=top_p,
-        frequency_penalty=frequency_penalty,
-        presence_penalty=presence_penalty
-    )
-    summarized_review = response.choices[0].message["content"]
+    chunked_reviews_JOIN = str("\n".join(chunked_reviews))
+    prompt = str(f"""
+    {summarize_prompt}
+
+    ```{str(chunked_reviews_JOIN)}```""")
+
+    response = genai_model.generate_content(prompt)
+    summarized_review = response.text
     return chunked_reviews, summarized_review
 
 
@@ -181,9 +171,10 @@ def main(
     # Check if necessary environment variables are set or not
     check_required_env_vars()
 
-    # Set the OpenAI API key
-    openai.base_url = "https://my-openai-gemini-demo.vercel.app"
-    openai.api_key = os.getenv("GEMINI_API_KEY")
+    # Set the Gemini API key
+    api_key = os.getenv("GEMINI_API_KEY")
+    genai.configure(api_key=api_key)
+    logger.debug(f"api_key: {api_key}")
 
     # Request a code review
     chunked_reviews, summarized_review = get_review(

requirements.txt

@@ -15,5 +15,4 @@ loguru==0.7.0
 
 requests==2.28.2
 
-langchain[openai]==0.0.144
-openai==0.27.0
+google-generativeai==0.5.2