Switch to using the more powerful GPT-3.5 LLM via OpenAI

README.md

@@ -30,7 +30,7 @@ cd AICodePal/src/python
 pip install -r requirements.txt
 ```
 
-3. Replace `YOUR_API_KEY` in the `ai_code_pal.py` file with your actual ChatGPT API key.
+3. Set the environment variable `OPENAI_API_KEY` with your actual ChatGPT API key.
 
 4. Now you can use the `get_chatgpt_suggestion` function from the `ai_code_pal.py` file in your project.
 

src/python/ai_code_pal.py

@@ -1,21 +1,25 @@
-import requests
-import os
+from typing import Tuple
+import openai
 
-def get_chatgpt_suggestion(code_snippet):
-    api_key = os.environ.get('OPENAI_API_KEY')
-    url = "https://api.openai.com/v1/engines/davinci-codex/completions"
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {api_key}"
-    }
-    data = {
-        "prompt": f"Improve the following Python code: {code_snippet}",
-        "max_tokens": 50,
-        "n": 1,
-        "stop": None,
-        "temperature": 0.7,
-    }
-    response = requests.post(url, headers=headers, json=data)
-    suggestion = response.json()['choices'][0]['text']
-    return suggestion.strip()
+
+OPENAI_MODEL = "gpt-3.5-turbo"
+
+
+def get_chatgpt_suggestion(code_snippet: str) -> Tuple[str, int]:
+    """
+    Returns PR style feedback for the given code snippet, and the number of tokens used.
+    """
+    system_prompt = f"""You are AI Code Pal. Suggest improvements to the provided Python code as in a code review."""
+
+    user_prompt = code_snippet
+    llm_response = openai.ChatCompletion.create(
+        model=OPENAI_MODEL,
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt},
+        ],
+        temperature=0.7,
+    )
+
+    return llm_response["choices"][0]["message"]["content"].strip(), llm_response["usage"]["total_tokens"]
 

src/python/requirements.txt

@@ -1,2 +1 @@
-requests
-
+openai