diff --git a/README.md b/README.md index de3e21f..e6eeb14 100644 --- a/README.md +++ b/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. diff --git a/src/python/ai_code_pal.py b/src/python/ai_code_pal.py index b2de671..70835cb 100644 --- a/src/python/ai_code_pal.py +++ b/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"] diff --git a/src/python/requirements.txt b/src/python/requirements.txt index 3288e92..ec838c5 100644 --- a/src/python/requirements.txt +++ b/src/python/requirements.txt @@ -1,2 +1 @@ -requests - +openai