feat: make model optional

codylabs · Sep 29, 2024 · d4d8e9c · d4d8e9c
1 parent f8ffca8
commit d4d8e9c
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/.github/workflows/code_review.yml/code_review.yml b/.github/workflows/code_review.yml/code_review.yml
@@ -26,5 +26,9 @@ jobs:
           pr_number: ${{ github.event.pull_request.number }}
           repository: ${{ github.repository }}
         env:
+          # This is provided by default, you do not need to change this
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          # Update the below in your GitHub repo actions secrets to use your OpenAI API key
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          # Update the below to use the model of your choice
+          OPENAI_MODEL: "gpt-4o"
diff --git a/src/config.py b/src/config.py
@@ -5,4 +5,5 @@
 
 OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
 GITHUB_TOKEN = os.getenv('GITHUB_TOKEN')
-DEBUG_MODE = os.getenv('DEBUG_MODE', 'false').lower() == 'true'
+DEBUG_MODE = os.getenv('DEBUG_MODE', 'false').lower() == 'true'
+OPENAI_MODEL = os.getenv('OPENAI_MODEL', 'gpt-4o')
diff --git a/src/model.py b/src/model.py
@@ -7,6 +7,7 @@
 logging.basicConfig(level=logging.DEBUG if config.DEBUG_MODE else logging.INFO)
 
 OPENAI_API_KEY = config.OPENAI_API_KEY
+OPENAI_MODEL = config.OPENAI_MODEL
 
 client = OpenAI(
     api_key=OPENAI_API_KEY,
@@ -18,7 +19,7 @@ def query_openai(prompt: str, retries=3, base_delay=1.0) -> str:
     while attempt < retries:
         try:
             completion = client.chat.completions.create(
-                model="gpt-4o",
+                model=OPENAI_MODEL,
                 messages=[
                     {"role": "system", "content": "You are a senior software engineer at Google reviewing a pull request."},
                     {"role": "user", "content": prompt}