Update app.py

app.py

@@ -1,131 +1,202 @@
 import streamlit as st
-import subprocess
 import os
-import time
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+import asyncio
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_chroma import Chroma
-from langchain.chains import RetrievalQA
-from langchain.memory import ConversationBufferMemory
 from langchain.prompts import PromptTemplate
+from langchain.memory import ConversationBufferMemory
 from langchain.schema import AIMessage, HumanMessage
 
-# Set OpenAI API Key
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-
-# Initialize OpenAI Chat Model
-llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.3, openai_api_key=OPENAI_API_KEY)
-
-# Initialize ChromaDB Vector Store
-vectorstore = Chroma(embedding_function=OpenAIEmbeddings())
-
-# Define Prompt Template
+# ------------------------------- 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
+
+# ------------------------------- CACHING -------------------------------
+@st.cache_resource
+def get_openai_model():
+    """Load OpenAI model once (cached) for efficiency."""
+    return ChatOpenAI(model_name="gpt-4-turbo", temperature=0.7)
+
+@st.cache_resource
+def get_vectorstore():
+    """Load ChromaDB once (cached) for faster retrieval."""
+    return Chroma(
+        embedding_function=OpenAIEmbeddings(),
+        persist_directory="/home/user/chroma_db"
+    )
+
+@st.cache_resource
+def get_memory():
+    """Initialize memory to store previous interactions."""
+    return ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+
+# Initialize models and memory
+llm = get_openai_model()
+vectorstore = get_vectorstore()
+memory = get_memory()
+
+# ------------------------------- PROMPT TEMPLATES -------------------------------
 prompt_template = PromptTemplate(
-    input_variables=["language", "description"],
+    input_variables=["language", "description", "context"],
     template="""
-    You are an **expert AI coding assistant** proficient in **ALL web development languages**, including:
-    - HTML, CSS, JavaScript, React, Vue.js, Angular, Svelte
-    - Node.js, Express.js, Django, Flask, Ruby on Rails
-    - PHP, Laravel, ASP.NET, .NET Core, C#
-    - Python, TypeScript, GraphQL, SQL, NoSQL
-    - WebSockets, REST APIs, WebAssembly, Web3.js
+    You are an **expert AI coding assistant** with memory.
+
+    **Context from previous interactions:**
+    {context}
 
-    **Task:** 
-    Generate a **complete, structured, and functional** code snippet based on the user request.
+    **Current Task:** Generate a structured and functional code snippet.
 
     **Language:** {language}  
     **Description:** {description}  
 
-    - Ensure the code follows best practices and industry standards.
+    - Ensure the code follows best practices.
     - Include comments explaining key sections.
     - If applicable, provide installation/setup instructions.
     """
 )
 
-# Streamlit UI
-st.set_page_config(page_title="Intelligent Coding Agent", layout="wide")
-st.title("🧠💻 Intelligent Coding Agent")
-st.markdown("### Generate and Run Code in Any Web Development Language!")
+enhance_template = PromptTemplate(
+    input_variables=["code", "context"],
+    template="""
+    Improve the given code based on context:
+
+    **Previous Conversation Context:**
+    {context}
+
+    **Original Code:**
+    {code}
+
+    **Enhanced Code with better efficiency, security, and structure:**
+    """
+)
+
+modify_template = PromptTemplate(
+    input_variables=["code", "modification", "context"],
+    template="""
+    Modify the code as per user requirements while maintaining context.
+
+    **Previous Conversation Context:**
+    {context}
+
+    **Original Code:**
+    {code}
+
+    **User Requested Modification:**
+    {modification}
+
+    **Modified Code:**
+    """
+)
+
+# ------------------------------- STREAMLIT UI -------------------------------
+st.title("🧠💻 Context-Aware Intelligent Coding Agent")
+st.markdown("### Generate, Enhance, and Modify Code While Retaining Context!")
 
-# Language selection
 languages = ["Python", "JavaScript", "PHP", "C#", "HTML/CSS/JS"]
 language = st.selectbox("Select a programming language:", languages)
-
-# User input for code description
 description = st.text_area("Describe the code you need:", placeholder="E.g., Build a login system in PHP")
 
-# Generate button
-if st.button("Generate Code"):
+# ------------------------------- SESSION STATE -------------------------------
+if "generated_code" not in st.session_state:
+    st.session_state.generated_code = ""
+
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = []
+
+# ------------------------------- ASYNC FUNCTIONS -------------------------------
+async def generate_code():
+    """Generates code while retaining context."""
     if not description.strip():
         st.warning("Please enter a code description!")
-    else:
-        # Show Progress Bar
-        progress_bar = st.progress(0)
-        status_text = st.empty()
-
-        # Progress Simulation
-        for i in range(50):
-            progress_bar.progress(i + 1)
-            status_text.text(f"🔄 Generating code... ({i+1}%)")
-            time.sleep(0.03)
-
-        # Generate Code
-        prompt = prompt_template.format(language=language, description=description)
-        response = llm.invoke([HumanMessage(content=prompt)])
-
-        # Progress Complete
-        for i in range(50, 100):
-            progress_bar.progress(i + 1)
-            status_text.text(f"✅ Code Generation Complete ({i+1}%)")
-            time.sleep(0.01)
-
-        # Show Generated Code
-        st.subheader("Generated Code:")
-        st.code(response.content, language=language.lower())
-
-        # Save to ChromaDB
-        vectorstore.add_texts([response.content])
-
-# Run Code Section
-st.subheader("Run Code (For Python, JS, PHP, C#)")
-if st.button("Run Code"):
-    if language in ["Python", "JavaScript", "PHP", "C#"]:
-        # Save code to temp file
-        filename = f"temp_code.{language.lower()}"
-        with open(filename, "w") as f:
-            f.write(response.content)
-
-        # Determine execution command
-        if language == "Python":
-            command = ["python", filename]
-        elif language == "JavaScript":
-            command = ["node", filename]
-        elif language == "PHP":
-            command = ["php", filename]
-        elif language == "C#":
-            command = ["dotnet-script", filename]
-        
-        # Run the code using subprocess
-        try:
-            output = subprocess.run(command, capture_output=True, text=True, timeout=10)
-            st.subheader("Output:")
-            st.text(output.stdout if output.stdout else "No output")
-            if output.stderr:
-                st.subheader("Errors:")
-                st.text(output.stderr)
-        except Exception as e:
-            st.error(f"Error executing code: {str(e)}")
-
-# Render HTML/CSS/JS
-if language == "HTML/CSS/JS":
-    st.subheader("Live Preview:")
-    st.components.v1.html(response.content, height=500)
-
-# Download Option
-if st.button("Download Code"):
-    with open("generated_code.txt", "w") as f:
-        f.write(response.content)
-    st.download_button("Download Generated Code", "generated_code.txt", "text/plain")
+        return
+    
+    st.info("⏳ Generating your code...")
+
+    # Retrieve past interactions for context
+    previous_context = "\n".join(st.session_state.chat_history)
+
+    # Format prompt
+    prompt = prompt_template.format(language=language, description=description, context=previous_context)
+    
+    # Generate code asynchronously
+    response = await asyncio.to_thread(llm.invoke, [HumanMessage(content=prompt)])
+
+    # Store generated code and update chat history
+    st.session_state.generated_code = response.content
+    st.session_state.chat_history.append(f"User asked: {description}\nAI Response:\n{response.content}")
+
+    # Display code
+    st.subheader("Generated Code:")
+    st.code(st.session_state.generated_code, language=language.lower())
+
+    # Save to vectorstore for future reference
+    vectorstore.add_texts([st.session_state.generated_code])
+
+async def enhance_code():
+    """Enhances the generated code while considering past context."""
+    if not st.session_state.generated_code:
+        st.warning("Generate code first before enhancing!")
+        return
+
+    st.info("⏳ Enhancing the code for better performance...")
+
+    previous_context = "\n".join(st.session_state.chat_history)
+
+    prompt = enhance_template.format(code=st.session_state.generated_code, context=previous_context)
+    response = await asyncio.to_thread(llm.invoke, [HumanMessage(content=prompt)])
+
+    st.session_state.generated_code = response.content
+    st.session_state.chat_history.append(f"AI Enhanced Code:\n{response.content}")
+
+    st.subheader("Enhanced Code:")
+    st.code(st.session_state.generated_code, language=language.lower())
+
+async def modify_code():
+    """Modifies existing code based on user input while remembering context."""
+    modification = st.text_area("Specify modifications (optional):", placeholder="E.g., Add validation for login form")
+
+    if st.button("Apply Modifications"):
+        if not st.session_state.generated_code:
+            st.warning("Generate code first before modifying!")
+            return
+        if not modification.strip():
+            st.warning("Please enter modification details!")
+            return
+
+        st.info("⏳ Modifying the code as per request...")
+
+        previous_context = "\n".join(st.session_state.chat_history)
+
+        prompt = modify_template.format(code=st.session_state.generated_code, modification=modification, context=previous_context)
+        response = await asyncio.to_thread(llm.invoke, [HumanMessage(content=prompt)])
+
+        st.session_state.generated_code = response.content
+        st.session_state.chat_history.append(f"User Modification Request: {modification}\nAI Modified Code:\n{response.content}")
+
+        st.subheader("Modified Code:")
+        st.code(st.session_state.generated_code, language=language.lower())
+
+# ------------------------------- BUTTONS -------------------------------
+col1, col2, col3 = st.columns(3)
+
+with col1:
+    if st.button("Generate Code"):
+        asyncio.run(generate_code())
+
+with col2:
+    if st.button("Enhance Code"):
+        asyncio.run(enhance_code())
+
+with col3:
+    modify_code()
+
+# ------------------------------- DOWNLOAD OPTION -------------------------------
+if st.session_state.generated_code:
+    st.download_button("Download Code", st.session_state.generated_code, "generated_code.txt", "text/plain")
+
+# ------------------------------- DISPLAY MEMORY -------------------------------
+st.subheader("Chat History & Context")
+st.text_area("Conversation Memory", value="\n".join(st.session_state.chat_history), height=200, disabled=True)
 
 st.markdown("---")
-st.markdown("🔹 **Built with LangChain (Latest), OpenAI GPT, and ChromaDB** 🔹")
+st.markdown("🔹 **Built with LangChain (Optimized), OpenAI GPT, and ChromaDB** 🔹")