Add python code to displayed answer

app.py

@@ -5,6 +5,19 @@ from langchain import OpenAI
 from langchain.chains import PALChain
 import datetime
 
+# Console to variable
+from io import StringIO
+import sys
+
+
+# tmp = sys.stdout
+# my_result = StringIO()
+# sys.stdout = my_result
+# print("HELLO");
+# sys.stdout = tmp
+# print(my_result.getvalue().lower())
+
+
 gpt_only_prompt = "Calculate the following, giving only the final answer:\n"
 
 
@@ -57,13 +70,29 @@ def calc_gpt_pal(math_problem, pal_chain):
     if not pal_chain:
         return "<pre>Please paste your OpenAI API key</pre>"
 
+    import re
+    tmp = sys.stdout
+    python_code_str_io = StringIO()
+    sys.stdout = python_code_str_io
     answer = pal_chain.run(math_problem)
+    sys.stdout = tmp
+    python_code = python_code_str_io.getvalue()
+
+    # Keep everything after the second set of """ in python_code
+    python_code = re.sub(r'.*"""', '', python_code, flags=re.DOTALL)
+
+    # Remove everything after the first [0m in python_code
+    python_code = re.sub(r'\[0m.*', '', python_code, flags=re.DOTALL)
+
+    # Remove leading and trailing whitespace from each line in python_code
+    python_code = os.linesep.join([s.strip() for s in python_code.splitlines()])
+
 
     print("\n==== date/time: " + str(datetime.datetime.now()) + " ====")
     print("calc_gpt_pal math problem: " + math_problem)
     print("calc_gpt_pal answer: " + answer)
 
-    html = "<pre>" + answer + "</pre>"
+    html = "<pre>" + python_code + "\n\n" + answer + "</pre>"
     return html