Update app.py

app.py

@@ -1,20 +1,163 @@
-from transformers import pipeline
+import os
+import time
 import streamlit as st
+import google.generativeai as genai
+from dotenv import load_dotenv
 
-# Load grammar correction pipeline
-@st.cache_resource
-def load_model():
-    return pipeline("text2text-generation", model="vennify/t5-base-grammar-correction")
+# Configuration
+load_dotenv()
+genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
+model = genai.GenerativeModel('gemini-2.0-flash')
 
-corrector = load_model()
+# Custom CSS for Chainlit-like appearance
+st.markdown("""
+<style>
+    /* Ensures proper contrast in dark mode */
+    body, .stApp {
+        background-color: #121212 !important;
+        color: #e0e0e0 !important;
+    }
+    /* Chat messages */
+    [data-testid="stChatMessage"] {
+        padding: 1rem;
+        border-radius: 15px;
+        margin: 1rem 0;
+        box-shadow: 0 2px 8px rgba(255,255,255,0.1);
+    }
+    /* User message */
+    .user-message {
+        background: #1e1e1e !important;
+        border: 1px solid #333 !important;
+        color: #ffffff !important;
+    }
+    /* Assistant message */
+    .assistant-message {
+        background: #2a2a2a !important;
+        border: 1px solid #444 !important;
+        color: #ffffff !important;
+    }
+    /* Chat input */
+    .stChatInput {
+        position: fixed;
+        bottom: 2rem;
+        width: 75%;
+        left: 50%;
+        transform: translateX(-50%);
+        background: #222 !important;
+        color: #fff !important;
+        border: 1px solid #555 !important;
+    }
+    /* Fix for buttons */
+    .stButton>button {
+        background-color: #333 !important;
+        color: white !important;
+        border: 1px solid #555 !important;
+    }
+</style>
+""", unsafe_allow_html=True)
 
-# Streamlit UI
-st.title("Grammar Correction Assistant")
+# App header
+st.title("📝 Grammar Guardian")
+st.caption("Hello, I'm the F&F Developers Grammatical Correction Agent. How may I help you today?")
 
-user_input = st.text_area("Enter a sentence to correct:", "She don't like going to the gym because it make her tired.")
+# Initialize session state
+if "history" not in st.session_state:
+    st.session_state.history = []
 
-if st.button("Correct Sentence"):
-    with st.spinner("Correcting..."):
-        result = corrector(user_input, max_length=100, clean_up_tokenization_spaces=True)
-        corrected_sentence = result[0]['generated_text']
-        st.markdown(f"**Corrected Sentence:** {corrected_sentence}")
+# Chat container
+chat_container = st.container()
+
+# Display history
+with chat_container:
+    for message in st.session_state.history:
+        if message["role"] == "user":
+            with st.chat_message("user", avatar="👤"):
+                st.markdown(f"**You**: {message['content']}")
+        else:
+            with st.chat_message("assistant", avatar="🤖"):
+                st.markdown(f"**Correction**: {message['content']['correction']}")
+                st.markdown(f"**Explanation**: {message['content']['explanation']}")
+                st.markdown(f"*<small>{message['content']['timestamp']}</small>*", 
+                          unsafe_allow_html=True)
+
+# Chat input
+prompt = st.chat_input("Type here!")
+
+if prompt:
+    # Add user message
+    with chat_container:
+        with st.chat_message("user", avatar="👤"):
+            st.markdown(f"**You**: {prompt}")
+    
+    st.session_state.history.append({"role": "user", "content": prompt})
+
+    # Generate response
+    with st.spinner("Analyzing..."):
+        try:
+            response = model.generate_content(
+                f"""You are an advanced grammar correction assistant. 
+                When a user provides a sentence, first correct the sentence and then 
+                explain why the correction was necessary. Highlight grammatical mistakes, 
+                sentence structure issues, and word choice errors in a simple and understandable way. 
+                Always provide detailed explanations to help the user learn.
+    
+                Please provide the response in the following structured format:
+                **Correction:** <Corrected Sentence>
+                **Explanation:** <Brief Explanation of Errors>
+    
+                Sentence: {prompt}
+                """
+            )
+            
+            # Ensure response exists
+            response_text = response.text if response and hasattr(response, "text") else "Error in response"
+    
+            # Extract correction and explanation
+            correction = ""
+            explanation = ""
+    
+            if "**Correction:**" in response_text and "**Explanation:**" in response_text:
+                parts = response_text.split("**Explanation:**")
+                correction = parts[0].replace("**Correction:**", "").strip()
+                explanation = parts[1].strip()
+            else:
+                # Fallback if formatting isn't followed
+                lines = response_text.split("\n")
+                if lines:
+                    correction = lines[0].strip()
+                    explanation = " ".join(lines[1:]).strip() if len(lines) > 1 else "Explanation not provided by the model."
+    
+            result = {
+                "correction": f"✨ {correction}",
+                "explanation": f"📚 {explanation}",
+                "timestamp": time.strftime("%H:%M:%S")
+            }
+            
+            # Add assistant message
+            with chat_container:
+                with st.chat_message("assistant", avatar="🤖"):
+                    st.markdown(f"**Correction**: {result['correction']}")
+                    st.markdown(f"**Explanation**: {result['explanation']}")
+                    st.markdown(f"*<small>{result['timestamp']}</small>*", unsafe_allow_html=True)
+            
+            st.session_state.history.append({
+                "role": "assistant",
+                "content": result
+            })
+    
+        except Exception as e:
+            st.error(f"Error processing request: {str(e)}")
+
+
+# Sidebar with history controls
+with st.sidebar:
+    st.header("Session Info")
+    st.write(f"Messages: {len(st.session_state.history)}")
+    
+    if st.button("Clear History"):
+        st.session_state.history = []
+        st.rerun()
+    
+    st.markdown("---")
+    st.markdown("**How to use:**")
+    st.markdown("1. Type your sentence\n2. Get instant feedback\n3. Review explanations")