diff --git a/ai_functions.py b/ai_functions.py
index a0d53a3..7558492 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,36 @@ 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']
+    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
+    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..0630c60
--- /dev/null
+++ b/test_code_optimizations.py
@@ -0,0 +1,28 @@
+from ai_functions import optimize_function
+import openai
+import keys
+
+# 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.
+        """
+        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
+
+
+# test_optimized_function_returns_correct_result("gpt-3.5-turbo")
+test_optimized_function_returns_correct_result("gpt-4")