Update app.py

app.py

@@ -6,7 +6,7 @@ from langchain_chroma import Chroma
 from langchain.prompts import PromptTemplate
 from langchain.memory import ConversationBufferMemory
 from langchain.schema import AIMessage, HumanMessage
-
+import re
 # ------------------------------- CONFIG -------------------------------
 st.set_page_config(page_title="Intelligent Coding Agent", layout="wide")
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")  # Ensure API key is set
@@ -103,8 +103,10 @@ if "generated_code" not in st.session_state:
     st.session_state.generated_code = ""
 
 # ------------------------------- ASYNC FUNCTIONS -------------------------------
+
+
 async def generate_code():
-    """Generates optimized code while keeping a conversational flow."""
+    """Generates optimized code while ensuring correct formatting."""
     if not description.strip():
         st.warning("Please enter a description!")
         return
@@ -117,8 +119,12 @@ async def generate_code():
     # Generate code asynchronously
     response = await asyncio.to_thread(llm.invoke, [HumanMessage(content=prompt)])
 
-    # Store generated code and maintain conversation flow
-    st.session_state.generated_code = response.content
+    # Extract only the code part (removes ``` and language tag)
+    extracted_code = re.sub(r"```[\w+]*\n", "", response.content)  # Remove triple backticks & language tag
+    extracted_code = re.sub(r"```$", "", extracted_code.strip())  # Remove trailing backticks
+
+    # Store formatted code in session state
+    st.session_state.generated_code = extracted_code
 
     # Display code in a properly formatted block
     st.subheader("Here’s your generated code:")
@@ -127,6 +133,7 @@ async def generate_code():
     # Save to vectorstore for reference
     vectorstore.add_texts([st.session_state.generated_code])
 
+
 async def enhance_code():
     """Enhances existing code with better efficiency, structure, and readability."""
     if not st.session_state.generated_code: