Update app.py

app.py

@@ -4,8 +4,9 @@ import os
 import json
 import base64
 from dotenv import load_dotenv
-# Correct import for the local storage library
 from streamlit_local_storage import LocalStorage
+from PIL import Image
+from io import BytesIO
 
 # --- PAGE CONFIGURATION ---
 st.set_page_config(
@@ -32,6 +33,58 @@ def convert_role_for_gemini(role):
         return "model"
     return role  # "user" stays the same
 
+def should_generate_visual(user_prompt, ai_response):
+    """Determine if a visual aid would be helpful based on the content"""
+    visual_keywords = [
+        'graph', 'plot', 'diagram', 'chart', 'visual', 'picture', 'illustration',
+        'geometry', 'triangle', 'circle', 'rectangle', 'square', 'polygon',
+        'coordinate', 'axis', 'function', 'parabola', 'line', 'slope',
+        'fraction', 'percentage', 'ratio', 'proportion', 'angles',
+        'number line', 'timeline', 'distribution', 'probability',
+        'pattern', 'sequence', 'series', 'matrix', 'vector'
+    ]
+    
+    combined_text = (user_prompt + " " + ai_response).lower()
+    return any(keyword in combined_text for keyword in visual_keywords)
+
+def generate_math_visual(user_prompt, ai_response):
+    """Generate a math visual using Gemini's image generation capabilities"""
+    try:
+        # Create a focused prompt for mathematical illustration
+        visual_prompt = f"""Create a clear, educational math diagram to illustrate this concept: {user_prompt}
+        
+        Style: Clean, minimalist educational diagram with:
+        - White or light background
+        - Clear labels and text
+        - Professional textbook style
+        - High contrast for readability
+        - Mathematical accuracy
+        - Simple, focused design
+        
+        Make it suitable for a math student to understand the concept better."""
+        
+        # Configure the image generation model
+        image_model = genai.GenerativeModel("gemini-2.5-flash-image-preview")
+        
+        response = image_model.generate_content(
+            visual_prompt,
+            generation_config=genai.GenerationConfig(
+                response_modalities=["TEXT", "IMAGE"]
+            )
+        )
+        
+        # Extract image from response
+        for part in response.candidates[0].content.parts:
+            if part.inline_data is not None:
+                image = Image.open(BytesIO(part.inline_data.data))
+                return image
+        
+        return None
+        
+    except Exception as e:
+        st.error(f"Could not generate visual: {e}")
+        return None
+
 # --- API KEY & MODEL CONFIGURATION ---
 load_dotenv()
 api_key = None
@@ -43,12 +96,20 @@ except (KeyError, FileNotFoundError):
 
 if api_key:
     genai.configure(api_key=api_key)
+    
+    # Main text model
     model = genai.GenerativeModel(
         model_name="gemini-2.5-flash-lite",
         system_instruction="""
         You are "Math Jegna", an AI specializing exclusively in mathematics.
         Your one and only function is to solve and explain math problems.
-        You are an AI math tutor that primarily uses the Professor B methodology developed by Everard Barrett. Use the best method for the situation. Create images, infographics to help whenever possible.  This methodology is designed to activate children's natural learning capacities and present mathematics as a contextual, developmental story that makes sense.
+        You are an AI math tutor that primarily uses the Professor B methodology developed by Everard Barrett. Use the best method for the situation. This methodology is designed to activate children's natural learning capacities and present mathematics as a contextual, developmental story that makes sense.
+
+        IMPORTANT: When explaining mathematical concepts that would benefit from visual aids, explicitly mention that a visual diagram will be provided to help illustrate the concept. Use phrases like:
+        - "Let me show you this with a diagram..."
+        - "A visual representation will help clarify this..."
+        - "I'll create an illustration to demonstrate..."
+        
 Core Philosophy and Principles
 1. Contextual Learning Approach
 Present math as a story: Every mathematical concept should be taught as part of a continuing narrative that builds connections between ideas
@@ -144,7 +205,6 @@ if "chats" not in st.session_state:
         else:
             raise ValueError("No shared chat")
     except (TypeError, ValueError, Exception):
-        # Use the correct method: localS.getItem()
         saved_data_json = localS.getItem("math_mentor_chats")
         if saved_data_json:
             saved_data = json.loads(saved_data_json)
@@ -230,11 +290,14 @@ for chat_key in list(st.session_state.chats.keys()):
 active_chat = st.session_state.chats[st.session_state.active_chat_key]
 
 st.title(f"Math Mentor: {st.session_state.active_chat_key} 🧠")
-st.write("Stuck on a math problem? Just type it below, and I'll walk you through it step-by-step!")
+st.write("Stuck on a math problem? Just type it below, and I'll walk you through it step-by-step with visual aids when helpful!")
 
 for message in active_chat:
     with st.chat_message(name=message["role"], avatar="🧑‍💻" if message["role"] == "user" else "🧠"):
         st.markdown(message["content"])
+        # Display stored images if they exist
+        if "image" in message:
+            st.image(message["image"], caption="Math Concept Illustration", use_column_width=True)
 
 if user_prompt := st.chat_input():
     active_chat.append({"role": "user", "content": user_prompt})
@@ -244,14 +307,31 @@ if user_prompt := st.chat_input():
     with st.chat_message("assistant", avatar="🧠"):
         with st.spinner("Math Mentor is thinking... 🤔"):
             try:
+                # Generate text response first
                 chat_session = model.start_chat(history=[
                     {'role': convert_role_for_gemini(msg['role']), 'parts': [msg['content']]} 
-                    for msg in active_chat[:-1]
+                    for msg in active_chat[:-1] if 'content' in msg
                 ])
                 response = chat_session.send_message(user_prompt)
                 ai_response_text = response.text
                 st.markdown(ai_response_text)
-                active_chat.append({"role": "assistant", "content": ai_response_text})
+                
+                # Store the text response
+                message_data = {"role": "assistant", "content": ai_response_text}
+                
+                # Check if we should generate a visual aid
+                if should_generate_visual(user_prompt, ai_response_text):
+                    with st.spinner("Creating visual aid... 🎨"):
+                        visual_image = generate_math_visual(user_prompt, ai_response_text)
+                        if visual_image:
+                            st.image(visual_image, caption="Math Concept Illustration", use_column_width=True)
+                            # Convert image to bytes for storage
+                            img_buffer = BytesIO()
+                            visual_image.save(img_buffer, format="PNG")
+                            img_bytes = img_buffer.getvalue()
+                            message_data["image"] = img_bytes
+                
+                active_chat.append(message_data)
 
             except Exception as e:
                 error_message = f"Sorry, something went wrong. Math Mentor is taking a break! 🤖\n\n**Error:** {e}"
@@ -259,9 +339,17 @@ if user_prompt := st.chat_input():
                 active_chat.append({"role": "assistant", "content": error_message})
 
 # --- SAVE DATA TO LOCAL STORAGE ---
+# Note: We'll only save text content to local storage, not images (too large)
+simplified_chats = {}
+for chat_key, messages in st.session_state.chats.items():
+    simplified_chats[chat_key] = []
+    for message in messages:
+        simplified_message = {"role": message["role"], "content": message["content"]}
+        # Don't save image data to local storage
+        simplified_chats[chat_key].append(simplified_message)
+
 data_to_save = {
-    "chats": st.session_state.chats,
+    "chats": simplified_chats,
     "active_chat_key": st.session_state.active_chat_key
 }
-# Use the correct method: localS.setItem()
 localS.setItem("math_mentor_chats", json.dumps(data_to_save))