Update app.py

app.py

@@ -38,7 +38,7 @@ from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationalRetrievalChain
 from templates import css, bot_template, user_template
 from io import BytesIO
-
+from contextlib import redirect_stdout
 
 # page config and sidebar declares up front allow all other functions to see global class variables
 st.set_page_config(page_title="Python AI Pair Programmer", layout="wide")
@@ -56,6 +56,7 @@ with col1:
 context = {}
 
 
+
 def create_file(filename, prompt, response, should_save=True):
     if not should_save:
         return
@@ -63,52 +64,36 @@ def create_file(filename, prompt, response, should_save=True):
     # Extract base filename without extension
     base_filename, ext = os.path.splitext(filename)
 
-    # Initialize the combined content
-    combined_content = ""
-
-    # Add Prompt with markdown title and emoji
-    combined_content += "# Prompt 📝\n" + prompt + "\n\n"
-
-    # Add Response with markdown title and emoji
+    combined_content = "# Prompt 📝\n" + prompt + "\n\n"
     combined_content += "# Response 💬\n" + response + "\n\n"
 
-    # Check for code blocks in the response
-    resources = re.findall(r"```([\s\S]*?)```", response)
+    # Extract Python code blocks from the response
+    resources = re.findall(r"```python([\s\S]*?)```", response)
     for resource in resources:
-        # Check if the resource contains Python code
-        if "python" in resource.lower():
-            # Remove the 'python' keyword from the code block
-            cleaned_code = re.sub(r'^\s*python', '', resource, flags=re.IGNORECASE | re.MULTILINE)
-            
-            # Add Code Results title with markdown and emoji
-            combined_content += "# Code Results 🚀\n"
-
-            # Redirect standard output to capture it
-            original_stdout = sys.stdout
-            sys.stdout = io.StringIO()
-            
-            # Execute the cleaned Python code within the context
-            try:
-                exec(cleaned_code, context)
-                code_output = sys.stdout.getvalue()
-                combined_content += f"```\n{code_output}\n```\n\n"
-                realtimeEvalResponse = "# Code Results 🚀\n" + "```" + code_output + "```\n\n"
-                st.write(realtimeEvalResponse)
-                
-            except Exception as e:
-                combined_content += f"```python\nError executing Python code: {e}\n```\n\n"
-            
-            # Restore the original standard output
-            sys.stdout = original_stdout
-        else:
-            # Add non-Python resources with markdown and emoji
-            combined_content += "# Resource 🛠️\n" + "```" + resource + "```\n\n"
+        combined_content += "# Code Results 🚀\n"
+
+        # Execute the Python code and capture output
+        with io.StringIO() as buf:
+            with redirect_stdout(buf):
+                try:
+                    # Remove the 'python' keyword from the code block
+                    cleaned_code = re.sub(r'^\s*python', '', resource, flags=re.IGNORECASE | re.MULTILINE)
+                    exec(cleaned_code.strip())
+                    code_output = buf.getvalue()
+                    combined_content += f"```\n{code_output}\n```\n\n"
+                except Exception as e:
+                    combined_content += f"```python\nError executing Python code: {e}\n```\n\n"
 
     # Save the combined content to a Markdown file
     if should_save:
-        with open(f"{base_filename}.md", 'w') as file:
+        with open(filename, 'w') as file:
             file.write(combined_content)
 
+        # Create a Base64 encoded link for the file
+        with open(filename, 'rb') as file:
+            encoded_file = base64.b64encode(file.read()).decode()
+            href = f'<a href="data:file/markdown;base64,{encoded_file}" download="{filename}">Download File 📄</a>'
+            st.markdown(href, unsafe_allow_html=True)
 
 
 # Read it aloud