Update app.py

app.py

@@ -1,10 +1,18 @@
-import streamlit as st
+# app.py - Complete Enhanced ICodeGuru Chatbot
 import os
 import json
-from typing import List, Optional
+import uuid
+import time
+import base64
+import datetime
+from typing import List, Optional, Dict, Any
+import streamlit as st
+import streamlit.components.v1 as components
 import nest_asyncio
+from dataclasses import dataclass, asdict
+from pathlib import Path
 
-# LangChain imports
+# LangChain imports (your teammate's backend)
 from langchain.vectorstores import Chroma
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.document_loaders import JSONLoader, DirectoryLoader
@@ -17,10 +25,15 @@ from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationalRetrievalChain
 
+# Enhanced components
+from components import render_response_box, render_enhanced_response_box
+from user_manager import UserManager, UserProfile
+from chat_manager import ChatManager, ChatSession
+
 # Apply asyncio patch for Streamlit compatibility
 nest_asyncio.apply()
 
-# --- CONFIGURATION ---
+# ========== Configuration ==========
 GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
 if not GROQ_API_KEY:
     st.error("⚠️ GROQ_API_KEY environment variable is not set!")
@@ -30,10 +43,121 @@ GROQ_MODEL = "llama3-8b-8192"
 EMBEDDING_MODEL = "all-MiniLM-L6-v2"
 CHROMA_PERSIST_DIR = "./chroma_db"
 DOCS_DIR = "./docs"
+USER_DATA_DIR = "./user_data"
+CHAT_DATA_DIR = "./chat_data"
+
+# Ensure directories exist
+for directory in [USER_DATA_DIR, CHAT_DATA_DIR, DOCS_DIR]:
+    Path(directory).mkdir(exist_ok=True)
+
+# ========== Page Configuration ==========
+st.set_page_config(
+    page_title="ICodeGuru AI Assistant", 
+    page_icon="🤖", 
+    layout="centered",
+    initial_sidebar_state="expanded"
+)
 
-class LangChainRAGSystem:
+# Load CSS with error handling
+try:
+    with open("style.css") as f:
+        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
+except FileNotFoundError:
+    st.warning("style.css file not found. Using default styling.")
+
+# ========== Initialize Managers ==========
+@st.cache_resource
+def get_user_manager():
+    return UserManager(USER_DATA_DIR)
+
+@st.cache_resource
+def get_chat_manager():
+    return ChatManager(CHAT_DATA_DIR)
+
+user_manager = get_user_manager()
+chat_manager = get_chat_manager()
+
+# ========== Logo Function ==========
+def get_base64_image(image_path):
+    try:
+        with open(image_path, "rb") as img_file:
+            return f"data:image/jpeg;base64,{base64.b64encode(img_file.read()).decode()}"
+    except FileNotFoundError:
+        return "data:image/svg+xml;base64,PHN2ZyB3aWR0aD0iNjAiIGhlaWdodD0iNjAiIHZpZXdCb3g9IjAgMCA2MCA2MCIgZmlsbD0ibm9uZSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIj4KPGNpcmNsZSBjeD0iMzAiIGN5PSIzMCIgcj0iMzAiIGZpbGw9IiM2NjdlZWEiLz4KPHR5cGUgPSJ0ZXh0Ij5JQzwvdGV4dD4KPC9zdmc+"
+
+# ========== User Authentication ==========
+def render_user_auth():
+    """Render user authentication interface"""
+    if 'user_id' not in st.session_state:
+        st.session_state.user_id = None
+    
+    if not st.session_state.user_id:
+        st.sidebar.markdown("### 👤 User Profile")
+        
+        auth_option = st.sidebar.radio("Choose option:", ["Login", "Create New Profile"])
+        
+        if auth_option == "Create New Profile":
+            with st.sidebar.form("create_profile"):
+                username = st.text_input("Username", placeholder="Enter username")
+                display_name = st.text_input("Display Name", placeholder="Your display name")
+                expertise_level = st.selectbox("Programming Experience", 
+                    ["Beginner", "Intermediate", "Advanced", "Expert"])
+                preferred_languages = st.multiselect("Preferred Languages", 
+                    ["Python", "JavaScript", "Java", "C++", "C#", "Go", "Rust", "PHP", "Ruby"])
+                learning_goals = st.text_area("Learning Goals", 
+                    placeholder="What do you want to learn?")
+                
+                if st.form_submit_button("Create Profile"):
+                    if username and display_name:
+                        try:
+                            profile = UserProfile(
+                                user_id=str(uuid.uuid4()),
+                                username=username,
+                                display_name=display_name,
+                                expertise_level=expertise_level,
+                                preferred_languages=preferred_languages,
+                                learning_goals=learning_goals
+                            )
+                            user_manager.create_user(profile)
+                            st.session_state.user_id = profile.user_id
+                            st.session_state.current_user = profile
+                            st.rerun()
+                        except Exception as e:
+                            st.error(f"Error creating profile: {str(e)}")
+                    else:
+                        st.error("Username and Display Name are required!")
+        
+        else:  # Login
+            existing_users = user_manager.get_all_usernames()
+            if existing_users:
+                selected_username = st.sidebar.selectbox("Select Username", existing_users)
+                
+                if st.sidebar.button("Login"):
+                    profile = user_manager.get_user_by_username(selected_username)
+                    if profile:
+                        st.session_state.user_id = profile.user_id
+                        st.session_state.current_user = profile
+                        st.rerun()
+            else:
+                st.sidebar.info("No existing profiles. Create a new one!")
+    
+    else:
+        # User is logged in
+        user = st.session_state.get('current_user')
+        if user:
+            st.sidebar.markdown(f"### 👋 Welcome, {user.display_name}!")
+            st.sidebar.markdown(f"**Level:** {user.expertise_level}")
+            
+            if st.sidebar.button("Logout"):
+                st.session_state.user_id = None
+                st.session_state.current_user = None
+                if 'current_session_id' in st.session_state:
+                    del st.session_state.current_session_id
+                st.rerun()
+
+# ========== Enhanced LangChain RAG System ==========
+class EnhancedLangChainRAGSystem:
     def __init__(self):
-        """Initialize the LangChain RAG system components."""
         self.embeddings = None
         self.vectorstore = None
         self.llm = None
@@ -47,14 +171,12 @@ class LangChainRAGSystem:
     
     def setup_components(self):
         """Setup all LangChain components."""
-        # Initialize embeddings
         self.embeddings = HuggingFaceEmbeddings(
             model_name=EMBEDDING_MODEL,
             model_kwargs={'device': 'cpu'},
             encode_kwargs={'normalize_embeddings': True}
         )
         
-        # Initialize LLM
         self.llm = ChatGroq(
             groq_api_key=GROQ_API_KEY,
             model_name=GROQ_MODEL,
@@ -62,10 +184,7 @@ class LangChainRAGSystem:
             max_tokens=1024
         )
         
-        # Load or create vectorstore
         self.load_vectorstore()
-        
-        # Setup retrieval chain
         self.setup_retrieval_chain()
     
     def load_vectorstore(self):
@@ -76,9 +195,7 @@ class LangChainRAGSystem:
                 embedding_function=self.embeddings,
                 collection_name="icodeguru_knowledge"
             )
-            st.info("✅ Loaded existing knowledge base.")
         except Exception as e:
-            st.warning(f"Creating new knowledge base: {e}")
             self.vectorstore = Chroma(
                 persist_directory=CHROMA_PERSIST_DIR,
                 embedding_function=self.embeddings,
@@ -86,72 +203,73 @@ class LangChainRAGSystem:
             )
     
     def setup_retrieval_chain(self):
-        """Setup the conversational retrieval chain."""
-        # Custom prompt template
-        prompt_template = """You are an expert assistant for iCodeGuru, a programming education platform. 
-        Use the following context to answer the user's question comprehensively and accurately.
-        Always provide relevant video links, website links, or resources when available in the context.
-        If you don't know the answer based on the context, say so clearly.
-    
-        Context: {context}
-    
-        Chat History: {chat_history}
-    
-        Human: {question}
-    
-        Assistant: I'll help you with that based on the iCodeGuru knowledge base.
-    
-        """
-        
-        PROMPT = PromptTemplate(
-            template=prompt_template,
-            input_variables=["context", "chat_history", "question"]
-        )
-        
-        # Always try to create retriever - let it handle empty collections gracefully
-        try:
-            # Create retriever
-            retriever = self.vectorstore.as_retriever(
-                search_type="similarity",
-                search_kwargs={"k": 4}  # Retrieve top 4 most relevant chunks
-            )
+        """Setup the conversational retrieval chain with personalization."""
+        def get_personalized_prompt():
+            user = st.session_state.get('current_user')
+            if user:
+                user_context = f"""
+            User Profile Context:
+            - Name: {user.display_name}
+            - Experience Level: {user.expertise_level}
+            - Preferred Languages: {', '.join(user.preferred_languages) if user.preferred_languages else 'None specified'}
+            - Learning Goals: {user.learning_goals or 'None specified'}
             
-            # Create conversational retrieval chain
-            self.retrieval_chain = ConversationalRetrievalChain.from_llm(
-                llm=self.llm,
-                retriever=retriever,
-                memory=self.memory,
-                combine_docs_chain_kwargs={"prompt": PROMPT},
-                return_source_documents=True,
-                verbose=True
+            Please tailor your response to match the user's experience level and preferences.
+            """
+                        else:
+                            user_context = "User profile not available. Provide general guidance."
+                        
+                        return f"""You are an expert assistant for iCodeGuru, a programming education platform. 
+            {user_context}
+            
+            Use the following context to answer the user's question comprehensively and accurately.
+            Always provide relevant video links, website links, or resources when available in the context.
+            If you don't know the answer based on the context, say so clearly.
+            
+            Context: {{context}}
+            Chat History: {{chat_history}}
+            Human: {{question}}"""
+
+            # Part 2 of Enhanced App - RAG System Implementation and Features
+
+            PROMPT = PromptTemplate(
+                template=get_personalized_prompt(),
+                input_variables=["context", "chat_history", "question"]
             )
-            st.success("✅ Retrieval chain setup successfully!")
             
-        except Exception as e:
-            st.warning(f"⚠️ Retrieval chain setup issue: {str(e)}")
-            self.retrieval_chain = None
+            try:
+                retriever = self.vectorstore.as_retriever(
+                    search_type="similarity",
+                    search_kwargs={"k": 4}
+                )
+                
+                self.retrieval_chain = ConversationalRetrievalChain.from_llm(
+                    llm=self.llm,
+                    retriever=retriever,
+                    memory=self.memory,
+                    combine_docs_chain_kwargs={"prompt": PROMPT},
+                    return_source_documents=True,
+                    verbose=False
+                )
+                
+            except Exception as e:
+                self.retrieval_chain = None
     
     def load_and_process_documents(self) -> List[Document]:
         """Load and process JSON documents from the docs directory."""
         documents = []
         
         if not os.path.exists(DOCS_DIR):
-            st.error(f"❌ Documents directory '{DOCS_DIR}' not found!")
             return documents
         
-        # Get all JSON files
         json_files = [f for f in os.listdir(DOCS_DIR) if f.endswith('.json')]
         
         if not json_files:
-            st.warning(f"⚠️ No JSON files found in '{DOCS_DIR}' directory!")
             return documents
         
-        st.info(f"📂 Found {len(json_files)} JSON files to process...")
-        
         for filename in json_files:
             file_path = os.path.join(DOCS_DIR, filename)
             try:
-                # Use JSONLoader with proper schema
                 loader = JSONLoader(
                     file_path=file_path,
                     jq_schema='.[]',
@@ -159,16 +277,13 @@ class LangChainRAGSystem:
                 )
                 file_docs = loader.load()
                 
-                # Add source metadata
                 for doc in file_docs:
                     doc.metadata['source_file'] = filename
                     doc.metadata['file_path'] = file_path
                 
                 documents.extend(file_docs)
-                st.success(f"✅ Loaded {len(file_docs)} documents from {filename}")
                 
             except Exception as e:
-                st.error(f"❌ Error loading {filename}: {str(e)}")
                 continue
         
         return documents
@@ -183,88 +298,56 @@ class LangChainRAGSystem:
         )
         
         chunks = text_splitter.split_documents(documents)
-        st.info(f"📄 Created {len(chunks)} document chunks")
         return chunks
     
     def clear_knowledge_base(self):
         """Clear the existing knowledge base."""
         try:
             if self.vectorstore:
-                # Delete the collection
                 self.vectorstore.delete_collection()
-                st.success("🗑️ Cleared existing knowledge base")
-                
-                # Recreate empty vectorstore
                 self.vectorstore = Chroma(
                     persist_directory=CHROMA_PERSIST_DIR,
                     embedding_function=self.embeddings,
                     collection_name="icodeguru_knowledge"
                 )
         except Exception as e:
-            st.error(f"❌ Error clearing knowledge base: {str(e)}")
+            pass
     
     def ingest_documents(self):
         """Complete document ingestion pipeline."""
-        with st.spinner("🔄 Loading documents..."):
-            # Load documents
-            documents = self.load_and_process_documents()
-            
-            if not documents:
-                st.error("❌ No documents loaded. Please check your docs folder.")
-                return False
+        documents = self.load_and_process_documents()
         
-        with st.spinner("✂️ Splitting documents into chunks..."):
-            # Split documents
-            chunks = self.split_documents(documents)
-            
-            if not chunks:
-                st.error("❌ No document chunks created.")
-                return False
+        if not documents:
+            return False
         
-        with st.spinner("🧠 Creating embeddings and storing in vector database..."):
-            try:
-                # Clear existing data
-                self.clear_knowledge_base()
-                
-                # Add chunks to vectorstore
-                self.vectorstore.add_documents(chunks)
-                
-                # Persist the vectorstore
-                self.vectorstore.persist()
-                
-                st.success(f"✅ Successfully ingested {len(chunks)} document chunks!")
-                
-                # Force recreate retrieval chain with new data
-                self.setup_retrieval_chain()
-                
-                # Verify the setup worked
-                try:
-                    doc_count = self.vectorstore._collection.count()
-                    st.info(f"📊 Knowledge base now contains {doc_count} documents")
-                except:
-                    st.info("📊 Knowledge base updated successfully")
-                
-                return True
-                
-            except Exception as e:
-                st.error(f"❌ Error during ingestion: {str(e)}")
-                return False
+        chunks = self.split_documents(documents)
+        
+        if not chunks:
+            return False
+        
+        try:
+            self.clear_knowledge_base()
+            self.vectorstore.add_documents(chunks)
+            self.vectorstore.persist()
+            self.setup_retrieval_chain()
+            return True
+            
+        except Exception as e:
+            return False
                 
     def get_answer(self, question: str) -> dict:
         """Get answer for a user question."""
         if not self.retrieval_chain:
             return {
-                "answer": "⚠️ Knowledge base is empty. Please refresh the knowledge base first.",
+                "answer": "⚠️ Knowledge base is initializing. Please try again in a moment.",
                 "source_documents": []
             }
         
         try:
-            # Check if vectorstore has documents before querying
             doc_count = 0
             try:
                 doc_count = self.vectorstore._collection.count()
             except:
-                # If count fails, try a simple similarity search to test
                 try:
                     test_results = self.vectorstore.similarity_search("test", k=1)
                     doc_count = len(test_results) if test_results else 0
@@ -273,136 +356,312 @@ class LangChainRAGSystem:
             
             if doc_count == 0:
                 return {
-                    "answer": "⚠️ No documents found in knowledge base. Please refresh the knowledge base first.",
+                    "answer": "I'm ready to help! However, I don't have any specific documents loaded in my knowledge base right now. I can still answer general programming questions based on my training. Feel free to ask anything!",
                     "source_documents": []
                 }
             
-            # Get response from the chain
             response = self.retrieval_chain({"question": question})
             return response
             
         except Exception as e:
             return {
-                "answer": f"❌ Error getting answer: {str(e)}",
+                "answer": f"I apologize, but I encountered an issue processing your question. Could you please try rephrasing it?",
                 "source_documents": []
             }
     
     def reset_conversation(self):
-        """Reset the conversation memory and UI chat history."""
+        """Reset the conversation memory."""
         self.memory.clear()
-        # Also clear Streamlit session state messages
-        if "messages" in st.session_state:
-            st.session_state.messages = []
-        st.success("🔄 Conversation history cleared!")
 
 # Initialize the RAG system
 @st.cache_resource
 def get_rag_system():
     """Cache the RAG system to avoid reinitialization."""
-    return LangChainRAGSystem()
+    return EnhancedLangChainRAGSystem()
 
-def main():
-    """Main Streamlit application."""
-    st.set_page_config(
-        page_title="EduBot for iCodeGuru",
-        page_icon="🎓",
-        layout="wide",
-        initial_sidebar_state="expanded"
-    )
-    
-    # Header
-    st.title("🎓 EduBot for @icodeguru0")
-    st.markdown("**Powered by LangChain** | Ask anything based on pre-loaded iCodeGuru knowledge.")
-    
-    # Initialize RAG system
-    rag_system = get_rag_system()
+# ========== Session Management ==========
+def initialize_chat_session():
+    """Initialize or load chat session"""
+    if 'current_session_id' not in st.session_state:
+        user_id = st.session_state.get('user_id')
+        if user_id:
+            session_id = chat_manager.create_session(user_id)
+            st.session_state.current_session_id = session_id
+            st.session_state.messages = []
+        else:
+            st.session_state.messages = []
+    else:
+        # Load existing session messages
+        session = chat_manager.get_session(st.session_state.current_session_id)
+        if session:
+            st.session_state.messages = []
+            for msg in session.messages:
+                st.session_state.messages.append({
+                    "role": msg.role,
+                    "content": msg.content,
+                    "message_id": msg.message_id,
+                    "rating": msg.rating,
+                    "is_bookmarked": msg.is_bookmarked,
+                    "source_documents": msg.source_documents
+                })
+
+# ========== Chat History Management ==========
+def render_chat_history_sidebar():
+    """Render chat history in sidebar"""
+    if st.session_state.get('user_id'):
+        user_sessions = chat_manager.get_user_sessions(st.session_state.user_id)
+        
+        if user_sessions:
+            st.sidebar.markdown("### 💬 Chat History")
+            
+            for session in user_sessions[:10]:  # Show last 10 sessions
+                session_title = session.title[:30] + "..." if len(session.title) > 30 else session.title
+                
+                col1, col2 = st.sidebar.columns([3, 1])
+                
+                with col1:
+                    if st.button(session_title, key=f"session_{session.session_id}"):
+                        st.session_state.current_session_id = session.session_id
+                        initialize_chat_session()
+                        st.rerun()
+                
+                with col2:
+                    if st.button("🗑️", key=f"delete_{session.session_id}", help="Delete session"):
+                        chat_manager.delete_session(session.session_id)
+                        if st.session_state.get('current_session_id') == session.session_id:
+                            del st.session_state.current_session_id
+                        st.rerun()
+
+# ========== Enhanced Sidebar Features ==========
+def render_enhanced_sidebar():
+    """Render enhanced sidebar with all features"""
+    # User Authentication
+    render_user_auth()
     
-    # Sidebar for admin functions
-    with st.sidebar:
-        st.header("⚙️ Admin Panel")
+    if st.session_state.get('user_id'):
+        # Chat History
+        render_chat_history_sidebar()
         
-        if st.button("🔄 Refresh Knowledge Base", type="primary"):
-            success = rag_system.ingest_documents()
-            if success:
-                st.balloons()
+        st.sidebar.markdown("---")
         
-        if st.button("🗑️ Clear Conversation"):
-            rag_system.reset_conversation()
-            st.rerun()  # Force UI refresh
+        # New Chat Button
+        if st.sidebar.button("🆕 New Chat", type="primary"):
+            user_id = st.session_state.user_id
+            session_id = chat_manager.create_session(user_id)
+            st.session_state.current_session_id = session_id
+            st.session_state.messages = []
+            get_rag_system().reset_conversation()
+            st.rerun()
         
-        st.markdown("---")
-        st.subheader("📊 System Info")
+        # Model Selection
+        st.sidebar.markdown("### 🧠 AI Settings")
+        model_options = ["llama3-8b-8192", "llama3-70b-8192"]
+        selected_model = st.sidebar.selectbox("Choose LLM Model", model_options, index=0)
         
-        # Show vectorstore stats
-        if rag_system.vectorstore:
-            try:
-                doc_count = rag_system.vectorstore._collection.count()
-                st.metric("Documents in KB", doc_count)
-            except:
-                st.metric("Documents in KB", "N/A")
+        if selected_model != GROQ_MODEL:
+            GROQ_MODEL = selected_model
+            get_rag_system().llm.model_name = selected_model
+        
+        # Knowledge Base Management
+        st.sidebar.markdown("### 📚 Knowledge Base")
+        if st.sidebar.button("🔄 Refresh Knowledge Base"):
+            with st.spinner("Refreshing knowledge base..."):
+                success = get_rag_system().ingest_documents()
+                if success:
+                    st.sidebar.success("✅ Knowledge base refreshed!")
+                else:
+                    st.sidebar.warning("⚠️ No documents found to load")
+        
+        # Export Chat History
+        st.sidebar.markdown("### 📤 Export")
+        if st.sidebar.button("📄 Export Chat History"):
+            if st.session_state.get('current_session_id'):
+                export_data = chat_manager.export_chat_history(
+                    st.session_state.user_id, 
+                    st.session_state.current_session_id
+                )
+                if export_data:
+                    st.sidebar.download_button(
+                        label="⬇️ Download JSON",
+                        data=json.dumps(export_data, indent=2),
+                        file_name=f"chat_export_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
+                        mime="application/json"
+                    )
+        
+        # User Statistics
+        st.sidebar.markdown("### 📊 Your Stats")
+        user_stats = user_manager.get_user_stats(st.session_state.user_id)
+        chat_stats = chat_manager.get_chat_statistics(st.session_state.user_id)
+        
+        col1, col2 = st.sidebar.columns(2)
+        with col1:
+            st.metric("Total Chats", chat_stats.get('total_sessions', 0))
+        with col2:
+            st.metric("Messages", chat_stats.get('total_messages', 0))
+        
+        st.sidebar.metric("Bookmarks", chat_stats.get('bookmarked_messages', 0))
+        
+        # Bookmarked Messages
+        bookmarked = chat_manager.get_bookmarked_messages(st.session_state.user_id)
+        if bookmarked:
+            st.sidebar.markdown("### 🔖 Bookmarked Responses")
+            for bookmark in bookmarked[:5]:  # Show 5 most recent
+                message_preview = bookmark['message']['content'][:50] + "..."
+                if st.sidebar.button(message_preview, key=f"bookmark_{bookmark['message']['message_id']}"):
+                    # Show full bookmarked message
+                    st.sidebar.write(bookmark['message']['content'])
+
+# ========== Message Rating Handler ==========
+def handle_component_value():
+    """Handle component interactions (ratings, bookmarks)"""
+    if 'component_value' in st.session_state and st.session_state.component_value:
+        data = st.session_state.component_value
+        
+        if data.get('action') == 'rate_message':
+            chat_manager.rate_message(
+                data['session_id'], 
+                data['message_id'], 
+                data['rating']
+            )
+        
+        elif data.get('action') == 'bookmark_message'):
+            chat_manager.bookmark_message(
+                data['session_id'], 
+                data['message_id'], 
+                data['is_bookmarked']
+            )
         
-        st.markdown("---")
-        st.caption("🧠 **ChromaDB** for vector storage")
-        st.caption("⚡ **Groq LLM** for answers")
-        st.caption("🔗 **LangChain** for orchestration")
+        # Clear the component value
+        st.session_state.component_value = None
+
+# ========== Main App Logic ==========
+def main():
+    """Main application logic"""
+    
+    # Handle component interactions
+    handle_component_value()
+    
+    # Display logo and header
+    image_data_url = get_base64_image("10001.jpeg")
+    st.markdown(f"""
+    <div class="custom-header">
+        <h1><img src="{image_data_url}" class="chatbot-logo" alt="Bot" /> ICodeGuru AI Assistant</h1>
+    </div>
+    """, unsafe_allow_html=True)
+    
+    # Render enhanced sidebar
+    render_enhanced_sidebar()
+    
+    # Initialize RAG system
+    rag_system = get_rag_system()
     
-    # Main chat interface
-    st.markdown("---")
+    # Check if user is logged in
+    if not st.session_state.get('user_id'):
+        st.info("👈 Please login or create a profile to start chatting!")
+        return
+    
+    # Initialize chat session
+    initialize_chat_session()
+    
+    # Generate response function
+    def generate_response(user_query):
+        """Generate AI response using LangChain system"""
+        if not user_query or not user_query.strip():
+            return "Please provide a valid question."
+        
+        try:
+            response = rag_system.get_answer(user_query)
+            answer = response.get("answer", "I apologize, but I couldn't generate a response. Please try again.")
+            
+            source_docs = response.get("source_documents", [])
+            if source_docs:
+                sources_text = "\n\n📚 **Sources:**\n"
+                for i, doc in enumerate(source_docs[:2], 1):
+                    source_file = doc.metadata.get('source_file', 'Unknown')
+                    content_preview = doc.page_content[:100] + "..." if len(doc.page_content) > 100 else doc.page_content
+                    sources_text += f"{i}. {source_file}: {content_preview}\n"
+                
+                answer += sources_text
+            
+            return answer, [doc.metadata.get('source_file', '') for doc in source_docs]
+            
+        except Exception as e:
+            return "I apologize, but I encountered an issue processing your question. Could you please try again.", []
     
-    # Initialize session state for chat history
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
+    # Display chat messages
+    for i, msg in enumerate(st.session_state.messages):
+        with st.chat_message(msg["role"]):
+            if msg["role"] == "assistant":
+                message_id = msg.get("message_id", f"msg-{i}")
+                session_id = st.session_state.get("current_session_id", "")
+                
+                render_enhanced_response_box(
+                    msg["content"], 
+                    message_id, 
+                    session_id,
+                    is_bookmarked=msg.get("is_bookmarked", False),
+                    rating=msg.get("rating"),
+                    show_actions=True
+                )
+            else:
+                st.markdown(msg["content"])
     
-    # Display chat history
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
-            if "sources" in message and message["sources"]:
-                with st.expander("📚 Sources"):
-                    for i, source in enumerate(message["sources"], 1):
-                        st.markdown(f"**Source {i}:** {source}")
+    # Chat input
+    prompt = st.chat_input("Type your message...")
     
-    # User input
-    if prompt := st.chat_input("💬 Ask your question about iCodeGuru..."):
-        # Add user message to chat history
-        st.session_state.messages.append({"role": "user", "content": prompt})
+    if prompt:
+        # Add user message to session
+        user_message_id = chat_manager.add_message(
+            st.session_state.current_session_id, 
+            "user", 
+            prompt
+        )
+        
+        # Add to session state
+        st.session_state.messages.append({
+            "role": "user", 
+            "content": prompt,
+            "message_id": user_message_id
+        })
         
-        # Display user message
         with st.chat_message("user"):
             st.markdown(prompt)
         
-        # Get assistant response
+        # Generate and display assistant response
         with st.chat_message("assistant"):
-            with st.spinner("🤔 Thinking..."):
-                response = rag_system.get_answer(prompt)
-                answer = response.get("answer", "No answer available.")
-                source_docs = response.get("source_documents", [])
-                
-                st.markdown(answer)
-                
-                # Show sources if available
-                if source_docs:
-                    sources = []
-                    for doc in source_docs[:3]:  # Show top 3 sources
-                        source = doc.metadata.get('source_file', 'Unknown source')
-                        content_preview = doc.page_content[:100] + "..." if len(doc.page_content) > 100 else doc.page_content
-                        sources.append(f"{source}: {content_preview}")
-                    
-                    if sources:
-                        with st.expander("📚 Sources"):
-                            for i, source in enumerate(sources, 1):
-                                st.markdown(f"**Source {i}:** {source}")
-                        
-                        # Add to session state with sources
-                        st.session_state.messages.append({
-                            "role": "assistant", 
-                            "content": answer,
-                            "sources": sources
-                        })
-                    else:
-                        st.session_state.messages.append({"role": "assistant", "content": answer})
-                else:
-                    st.session_state.messages.append({"role": "assistant", "content": answer})
+            with st.spinner("Thinking..."):
+                full_response, source_docs = generate_response(prompt)
+            
+            # Add assistant message to session
+            assistant_message_id = chat_manager.add_message(
+                st.session_state.current_session_id, 
+                "assistant", 
+                full_response,
+                source_docs
+            )
+            
+            # Display response with enhanced box
+            render_enhanced_response_box(
+                full_response, 
+                assistant_message_id, 
+                st.session_state.current_session_id,
+                is_bookmarked=False,
+                rating=None,
+                show_actions=True
+            )
+            
+            # Add to session state
+            st.session_state.messages.append({
+                "role": "assistant", 
+                "content": full_response,
+                "message_id": assistant_message_id,
+                "rating": None,
+                "is_bookmarked": False,
+                "source_documents": source_docs
+            })
+        
+        # Update user chat count
+        user_manager.increment_chat_count(st.session_state.user_id)
 
 if __name__ == "__main__":
     main()