Update app.py

app.py

@@ -34,6 +34,11 @@ def generate_code_from_prompt(prompt, model="gpt-4o-mini"):
     )
     return response['choices'][0]['message']['content']
 
+# Function to save code to a .txt file
+def save_code_to_file(code, filename="generated_code.txt"):
+    with open(filename, "w") as f:
+        f.write(code)
+
 # Streamlit app starts here
 st.title("AI Assistance")
 
@@ -120,8 +125,25 @@ if openai_api_key:
                     # Generate code using GPT-4
                     generated_code = generate_code_from_prompt(code_prompt)
                     
+                    # Clean the generated code to ensure only code is saved (removing comments or additional text)
+                    clean_code = "\n".join([line for line in generated_code.splitlines() if not line.strip().startswith("#")])
+
+                    # Save the clean code to a file
+                    save_code_to_file(clean_code)
+
                     # Display the generated code
                     st.write("### Generated Code:")
-                    st.code(generated_code, language="python")
+                    st.code(clean_code, language="python")
+
+                    # Provide a download link for the generated code
+                    with open("generated_code.txt", "w") as f:
+                        f.write(clean_code)
+
+                    st.download_button(
+                        label="Download Generated Code",
+                        data=open("generated_code.txt", "rb").read(),
+                        file_name="generated_code.txt",
+                        mime="text/plain"
+                    )
             else:
                 st.error("Please provide a prompt to generate the code.")