Implements Code Improvement, powered by AI Functions.

scripts/ai_functions.py

@@ -14,3 +14,35 @@ def call_ai_function(function, args, description, model = "gpt-4"):
     )
 
     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: Optional[str] = None) -> str:
+    function_string = "def create_test_cases(code: str, focus: Optional[str] = None) -> str:"
+    args = [focus] if focus else []
+    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