From ab2eab50e2b9f678932feb6877a3b0da445bc62c Mon Sep 17 00:00:00 2001
From: Himanshu Tomar <himanshutomar3031@gmail.com>
Date: Wed, 26 Apr 2023 12:18:08 +0530
Subject: [PATCH 1/2] Implemented a AI function to enhance the time and space
 efficiency of another function.

---
 ai_functions.py            | 36 ++++++++++++++++++++++++++++++++++--
 test_code_optimizations.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 2 deletions(-)
 create mode 100644 test_code_optimizations.py

diff --git a/ai_functions.py b/ai_functions.py
index a0d53a3..cbc8444 100644
--- a/ai_functions.py
+++ b/ai_functions.py
@@ -1,9 +1,10 @@
 import openai
 
-def ai_function(function, args, description, model = "gpt-4"):
+
+def ai_function(function, args, description, model="gpt-4"):
     # parse args to comma separated 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. Do not include any other explanatory text in your response."},{"role": "user", "content": args}]
+    messages = [{"role": "system", "content": f"You are now the following python function: ```# {description}\n{function}```\n\nOnly respond with your `return` value. Do not include any other explanatory text in your response."}, {"role": "user", "content": args}]
 
     response = openai.ChatCompletion.create(
         model=model,
@@ -12,3 +13,34 @@ def ai_function(function, args, description, model = "gpt-4"):
     )
 
     return response.choices[0].message["content"]
+
+
+def optimize_function(function, args, model="gpt-4"):
+    # Generate a prompt to optimize the function
+    prompt = f"Optimize the following Python function:\n\n{function.__name__}({', '.join(args)})\n\nThe function's purpose is to {function.__doc__}\n\noptimize the function's time and space complexity.\n\nDo not include any other explanatory text in your response."
+    messages = [
+        {"role": "system",
+         "content": "You are a code optimizer and refactor. Just give only output not extra text with that."},
+        {"role": "user", "content": prompt}
+    ]
+
+    response = openai.ChatCompletion.create(
+        model=model,
+        messages=messages,
+        n=1,
+        stop=None,
+        temperature=0.5,
+    )
+
+    # Extract the suggested optimizations from the GPT response
+    optimized_code = response.choices[0]['message']['content']
+
+    # Compile the optimized function code into a function object
+    optimized_function = None
+    exec(optimized_code, globals(), locals())
+    for var in locals():
+        if var != '__builtins__':
+            optimized_function = locals()[var]
+            break
+
+    return optimized_function
diff --git a/test_code_optimizations.py b/test_code_optimizations.py
new file mode 100644
index 0000000..9760d84
--- /dev/null
+++ b/test_code_optimizations.py
@@ -0,0 +1,28 @@
+from ai_functions import optimize_function
+import openai
+import keys
+import inspect
+
+
+# Initialize the OpenAI API client
+openai.api_key = keys.OPENAI_API_KEY
+
+
+def test_optimized_function_returns_correct_result(model):
+    def calculate_square(x):
+        """
+        This function returns the square of its input.
+        """
+        return x**2
+
+    optimized_function = optimize_function(calculate_square, ["x"], model)
+    assert optimized_function(2) == 4
+    assert optimized_function(3) == 9
+    assert optimized_function(4) == 16
+    optimized_function_source = inspect.getsource(optimized_function)
+    print(
+        f"Here is the optimized function in terms of time and space complexity:\n{optimized_function_source}")
+
+
+# test_optimized_function_returns_correct_result("gpt-3.5-turbo")
+test_optimized_function_returns_correct_result("gpt-4")

From 667ac1b3d62e3bc5d896cbae64eb61cc7654a067 Mon Sep 17 00:00:00 2001
From: Himanshu Tomar <himanshutomar3031@gmail.com>
Date: Wed, 26 Apr 2023 16:24:20 +0530
Subject: [PATCH 2/2] Added better test case to understand the use case

---
 ai_functions.py            |  2 ++
 test_code_optimizations.py | 12 ++++++------
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/ai_functions.py b/ai_functions.py
index cbc8444..7558492 100644
--- a/ai_functions.py
+++ b/ai_functions.py
@@ -34,6 +34,8 @@ def optimize_function(function, args, model="gpt-4"):
 
     # Extract the suggested optimizations from the GPT response
     optimized_code = response.choices[0]['message']['content']
+    print(
+        f"Here is the optimized function in terms of time and space complexity:\n{optimized_code}")
 
     # Compile the optimized function code into a function object
     optimized_function = None
diff --git a/test_code_optimizations.py b/test_code_optimizations.py
index 9760d84..0630c60 100644
--- a/test_code_optimizations.py
+++ b/test_code_optimizations.py
@@ -1,8 +1,6 @@
 from ai_functions import optimize_function
 import openai
 import keys
-import inspect
-
 
 # Initialize the OpenAI API client
 openai.api_key = keys.OPENAI_API_KEY
@@ -13,15 +11,17 @@ def calculate_square(x):
         """
         This function returns the square of its input.
         """
-        return x**2
+        n = x
+        res = 0
+        while n > 0:
+            res += x
+            n -= 1
+        return res
 
     optimized_function = optimize_function(calculate_square, ["x"], model)
     assert optimized_function(2) == 4
     assert optimized_function(3) == 9
     assert optimized_function(4) == 16
-    optimized_function_source = inspect.getsource(optimized_function)
-    print(
-        f"Here is the optimized function in terms of time and space complexity:\n{optimized_function_source}")
 
 
 # test_optimized_function_returns_correct_result("gpt-3.5-turbo")