54 lines
1.8 KiB
Python
54 lines
1.8 KiB
Python
from typing import List, Optional
|
|
import json
|
|
import openai
|
|
|
|
|
|
def call_ai_function(function, args, description, model="gpt-4"):
|
|
# parse args to comma seperated string
|
|
args = ", ".join(args)
|
|
messages = [{"role": "system",
|
|
"content": f"You are now the following python function: ```# {description}\n{function}```\n\nOnly respond with your `return` value."},
|
|
{"role": "user",
|
|
"content": args}]
|
|
|
|
response = openai.ChatCompletion.create(
|
|
model=model,
|
|
messages=messages,
|
|
temperature=0
|
|
)
|
|
|
|
return response.choices[0].message["content"]
|
|
|
|
# Evaluating code
|
|
|
|
|
|
def evaluate_code(code: str) -> List[str]:
|
|
function_string = "def analyze_code(code: str) -> List[str]:"
|
|
args = [code]
|
|
description_string = """Analyzes the given code and returns a list of suggestions for improvements."""
|
|
|
|
result_string = call_ai_function(function_string, args, description_string)
|
|
return json.loads(result_string)
|
|
|
|
|
|
# Improving code
|
|
|
|
def improve_code(suggestions: List[str], code: str) -> str:
|
|
function_string = "def generate_improved_code(suggestions: List[str], code: str) -> str:"
|
|
args = [json.dumps(suggestions), code]
|
|
description_string = """Improves the provided code based on the suggestions provided, making no other changes."""
|
|
|
|
result_string = call_ai_function(function_string, args, description_string)
|
|
return result_string
|
|
|
|
|
|
# Writing tests
|
|
|
|
def write_tests(code: str, focus: List[str]) -> str:
|
|
function_string = "def create_test_cases(code: str, focus: Optional[str] = None) -> str:"
|
|
args = [code, json.dumps(focus)]
|
|
description_string = """Generates test cases for the existing code, focusing on specific areas if required."""
|
|
|
|
result_string = call_ai_function(function_string, args, description_string)
|
|
return result_string
|