Update app.py

app.py

@@ -2,292 +2,312 @@ import gradio as gr
 import io
 import base64
 import os
+import json
+import re
 from PIL import Image
 from huggingface_hub import InferenceClient
-from openai import OpenAI
+from google.generativeai import configure, GenerativeModel
+from google.ai.generativelanguage import Content, Part
 
 # Load API keys from environment variables
 inference_api_key = os.environ.get("HF_TOKEN")
-chat_api_key = os.environ.get("HF_TOKEN")
-chat_api_key2 = os.environ.get("OPENROUTER_TOKEN")
+google_api_key = os.environ.get("GOOGLE_API_KEY")  # New Google API key
 
+# Configure Google API
+configure(api_key=google_api_key)
 
-
-# Global variable to store the image data URL and prompt for the currently generated image.
+# Global variables to store the image data URL and prompt for the currently generated image.
 global_image_data_url = None
 global_image_prompt = None  # Still stored if needed elsewhere
 
-def generate_prompt_from_options(difficulty, age, level, extra_details=""):
+def update_difficulty_label(active_session):
+    return f"**Current Difficulty:** {active_session.get('difficulty', 'Very Simple')}"
+
+def generate_prompt_from_options(difficulty, age, autism_level, topic_focus, treatment_plan=""):
     """
-    Uses the OpenAI chat model (via Hugging Face Inference API) to generate an image generation prompt
-    based on the selected difficulty, age, autism level, and any extra details the user provides.
+    Generate an image prompt using Google's Gemini model.
     """
     query = (
-
         f"""
-        Follow the instructions below to Generate an image generation prompt for an educational image intended for Autistic children.
-        Consider the following parameters:\n
-          - Difficulty: {difficulty}\n
-          - Age: {age}\n
-          - Autism Level: {level}\n
-          - Extra Details: {extra_details}\n\n
-        Use the following system prompt to guide the image generation process:\n
-
-        System Prompt:
-
-        You are an image generation engine specializing in creating clear, calming, and visually supportive images designed for children with autism spectrum disorder (ASD). Your primary goal is to produce images that aid in understanding, communication, emotional regulation, and daily routines.  Prioritize the following characteristics:
-
-        **1. Clarity and Simplicity:**
-
-        # *   **Minimalist Backgrounds:**  Use solid, muted colors (e.g., soft blues, greens, light grays, pastels) or very simple, uncluttered backgrounds. Avoid busy patterns, highly contrasting colors, or distracting elements.
-        *   **Clear Subject Focus:**  The main subject of the image should be prominent and easily identifiable. Avoid unnecessary details that could cause confusion or sensory overload.
-        *   **Unambiguous Representations:**  Objects and people should be depicted in a realistic and straightforward manner. Avoid abstract art or overly stylized representations.  If depicting emotions, make them very clear and easily recognizable (e.g., a simple, wide smile for happiness, a single tear for sadness).
-
-        **2. Visual Structure and Predictability:**
-
-        *   **Literal Interpretation:**  The images should be highly literal.  Avoid metaphors, symbolism, or implied meanings. If depicting a sequence of events, make each step visually distinct.
-        *   **Defined Borders:**  Consider using clear outlines or borders around objects and people to enhance visual separation and definition.
-        *   **Consistent Style:**  Maintain a consistent visual style across multiple images.  This helps build familiarity and predictability.
-
-        **3. Sensory Considerations:**
-
-        *   **Soft Color Palette:**  Favor muted, calming colors. Avoid overly bright, saturated, or fluorescent colors.
-        *   **Reduced Visual Complexity:**  Limit the number of elements in the image to prevent sensory overload.
-        *   **Smooth Textures:**  If textures are depicted, they should appear smooth and non-threatening. Avoid rough, jagged, or overly detailed textures.
-
-        **4. Positive and Supportive Imagery:**
-
-        *   **Positive Reinforcement:**  Images should be encouraging and positive.  Depict success, cooperation, and positive social interactions.
-        *   **Calm and Relaxing Scenes:**  Consider scenes that promote calmness, such as nature scenes (e.g., a quiet forest, a calm beach), or familiar, safe environments (e.g., a cozy bedroom, a well-organized classroom).
-        *   **Avoidance of Triggers:**  Be mindful of potential triggers for anxiety or distress. Avoid images that depict conflict, overwhelming crowds, or potentially frightening situations.
-
-        **5. Specific Use Cases (Adapt as needed):**
-
-        *   **Social Stories:**  If generating images for a social story, ensure each image clearly illustrates a single step in the sequence.  Use consistent characters and settings throughout the story.
-        *   **Visual Schedules:**  If creating images for a visual schedule, make each activity easily identifiable and visually distinct.
-        *   **Emotion Recognition:**  If depicting emotions, use clear facial expressions and body language. Consider using a consistent character to represent different emotions.
-        *   **Communication Aids:**  If creating images for communication, ensure the objects or actions are clearly depicted and easily recognizable.
-        *  **Daily Routines**: Brushing teeth, eating food, going to school.
-        * **Learning concepts**: Shapes, colors, animals, numbers, alphabet.
-
-        **Prompting Instructions:**
-
-        When providing a prompt to the model, be as specific as possible, including:
-
-        *   **The subject of the image:**  "A boy brushing his teeth."
-        *   **The desired style:**  "Simple, clear, with a solid light blue background."
-        *   **The intended use:**  "For a visual schedule."
-        *   **Any specific details:**  "The boy should be smiling. The toothbrush should be blue."
-        *   **Emotions:** Clearly state the emotion "happy" or "calm."
-
-        **Example Prompts (using the above system prompt as a base):**
-
-        *   "Generate an image for a visual schedule. The subject is 'eating lunch.' Show a child sitting at a table with a plate of food (sandwich, apple slices, and a glass of milk). The background should be a solid, pale green. The child should be smiling. Use a clear, simple style with defined outlines."
-        *   "Generate an image to help with emotion recognition. The subject is 'sad.' Show a child's face with a single tear rolling down their cheek and a downturned mouth. The background should be a solid, light gray. Use a simple, realistic style."
-        *   "Generate an image for a social story about going to the doctor. Show a child sitting in a doctor's waiting room, calmly looking at a book. The room should have a few simple toys and a window. The background should be a soft blue. The style should be clear and uncluttered."
-        * "Generate a picture of two block shapes in a simple, cartoon style. One red square and one blue circle. Place them on a white background."
-        * "Generate a cartoon image of a dog. Make the dog appear to be friendly and non-threatening. Use warm colors."
-
-        Ensure your Prompts are acccurate and ensure the images are accurate and dont have any irregularities or deforamtions in them.
-        use descriptive and detailed prompts
+        Follow the instructions below to generate an image generation prompt for an educational image intended for autistic children.
+        Consider the following parameters:
+          - Difficulty: {difficulty}
+          - Age: {age}
+          - Autism Level: {autism_level}
+          - Topic Focus: {topic_focus}
+          - Treatment Plan: {treatment_plan}
+
+        Emphasize that the image should be clear, calming, and support understanding and communication. The style should match the difficulty level: for example, "Very Simple" produces very basic visuals while "Very Detailed" produces rich visuals.
+
+        The image should specifically focus on the topic: "{topic_focus}".
+
+        Please generate a prompt that instructs the image generation engine to produce an image with:
+        1. Clarity and simplicity (minimalist backgrounds, clear subject)
+        2. Literal representation with defined borders and consistent style
+        3. Soft, muted colors and reduced visual complexity
+        4. Positive, calm scenes
+        5. Clear focus on the specified topic
+
+        Use descriptive and detailed language.
         """
     )
 
-    messages = [
-        {
-            "role": "user",
-            "content": query
-        }
-    ]
-
-    client = OpenAI(
-        base_url="https://openrouter.ai/api/v1",
-        api_key=chat_api_key2
-    )
+    # Initialize the Gemini Pro model
+    model = GenerativeModel('gemini-2.0-flash-lite')
 
-    stream = client.chat.completions.create(
-        model="sophosympatheia/rogue-rose-103b-v0.2:free",
-        temperature=0.5,
-        messages=messages,
-        max_tokens=8192,
-        stream=True
-    )
+    # Generate content using the Gemini model
+    response = model.generate_content(query)
 
-    response_text = ""
-    for chunk in stream:
-        response_text += chunk.choices[0].delta.content
-    return response_text.strip()
+    return response.text.strip()
 
-def generate_image_fn(selected_prompt, guidance_scale=7.5, negative_prompt="ugly, blurry, poorly drawn hands , Lewd , nude", num_inference_steps=50):
+def generate_image_fn(selected_prompt, guidance_scale=7.5,
+                      negative_prompt="ugly, blurry, poorly drawn hands, lewd, nude, deformed, missing limbs, missing eyes, missing arms, missing legs",
+                      num_inference_steps=50):
     """
-    Uses the Hugging Face Inference API to generate an image from the provided prompt.
-    Converts the image to a data URL for later use and stores the prompt globally.
-
-    Additional parameters:
-      - guidance_scale: Influences how strongly the image generation adheres to the prompt.
-      - negative_prompt: Specifies undesirable elements to avoid in the generated image.
-      - num_inference_steps: The number of denoising steps for image generation.
+    Generate an image from the prompt via the Hugging Face Inference API.
+    Convert the image to a data URL.
     """
     global global_image_data_url, global_image_prompt
-
-    # Save the chosen prompt for potential future use.
     global_image_prompt = selected_prompt
-
-    image_client = InferenceClient(
-        provider="hf-inference",
-        api_key=inference_api_key
-    )
-
+    image_client = InferenceClient(provider="hf-inference", api_key=inference_api_key)
     image = image_client.text_to_image(
         selected_prompt,
-        model="stabilityai/stable-diffusion-3.5-large-turbo", #||           black-forest-labs/FLUX.1-dev  ||        stabilityai/stable-diffusion-3.5-large-turbo
+        model="stabilityai/stable-diffusion-3.5-large-turbo",
         guidance_scale=guidance_scale,
         negative_prompt=negative_prompt,
         num_inference_steps=num_inference_steps
     )
-
     buffered = io.BytesIO()
     image.save(buffered, format="PNG")
     img_bytes = buffered.getvalue()
     img_b64 = base64.b64encode(img_bytes).decode("utf-8")
     global_image_data_url = f"data:image/png;base64,{img_b64}"
-
     return image
 
-def generate_image_and_reset_chat(difficulty, age, level, extra_details, active_session, saved_sessions):
+def generate_image_and_reset_chat(age, autism_level, topic_focus, treatment_plan, active_session, saved_sessions):
     """
-    Saves any current active session into the saved sessions list. Then, using the three selected options and extra details,
-    generates an image generation prompt, creates an image, and starts a new active session.
+    Generate a new image (with the current difficulty) and reset the chat.
+    Now includes the topic_focus parameter to specify what the image should focus on.
     """
     new_sessions = saved_sessions.copy()
     if active_session.get("prompt"):
         new_sessions.append(active_session)
-
-    generated_prompt = generate_prompt_from_options(difficulty, age, level, extra_details)
-    image = generate_image_fn(generated_prompt)  # Uses default guidance_scale, negative_prompt, and num_inference_steps
-
-    new_active_session = {"prompt": generated_prompt, "image": global_image_data_url, "chat": []}
+    # Use the current difficulty from the active session (which should be updated if advanced)
+    current_difficulty = active_session.get("difficulty", "Very Simple")
+    generated_prompt = generate_prompt_from_options(current_difficulty, age, autism_level, topic_focus, treatment_plan)
+    image = generate_image_fn(generated_prompt)
+    new_active_session = {
+        "prompt": generated_prompt,
+        "image": global_image_data_url,
+        "chat": [],
+        "treatment_plan": treatment_plan,
+        "topic_focus": topic_focus,
+        "identified_details": [],
+        "difficulty": current_difficulty,
+        "autism_level": autism_level,
+        "age": age
+    }
     return image, new_active_session, new_sessions
 
-def compare_details_chat_fn(user_details):
+def compare_details_chat_fn(user_details, treatment_plan, chat_history, identified_details):
     """
-    Uses the vision language model to evaluate the user description based solely on the generated image.
-    The message includes both the image (using its data URL) and the user’s text.
+    Evaluate the child's description using Google's Gemini Vision model.
     """
     if not global_image_data_url:
         return "Please generate an image first."
 
-    # Prepare the message content as a list of parts:
-    # 1. The image part – here we send the image data URL (in practice, you might need to supply a public URL).
-    # 2. The text part – containing the user's description.
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {
-                    "type": "image_url",
-                    "image_url": {"url": global_image_data_url}
-                },
-                {
-                    "type": "text",
-                    "text": (
-                        f"Based on the image provided above, please evaluate the following description given by the child:\n"
-                        f"'{user_details}'\n\n"
-                        """
-                        You are a friendly and encouraging teacher, guiding a child in describing an image. Speak directly to the child using simple, clear language. Provide positive reinforcement when the child gives a correct or accurate description.
-
-                        If the child's description is incorrect or inaccurate, gently guide them with hints rather than direct corrections. Use Hint before providing guidance. Keep your hints playful and engaging to encourage curiosity.
-
-                        Avoid repeating the child’s description. Instead, focus on giving feedback based on the image. If the description is correct, acknowledge it warmly with praise.
-
-                        Keep the conversation going by asking open-ended questions about the image to encourage the child to observe and think more deeply. Use questions that spark curiosity, such as 'What else do you see?' or 'Why do you think that is happening?'
-
-                        Do not mention your own thoughts, system prompts, or provide direct details about the image. Stay fully engaged in a natural, conversational way, making learning fun and interactive!
-
-                        """
-                    )
-                }
-            ]
-        }
-    ]
-
-    chat_client = OpenAI(
-        base_url="https://openrouter.ai/api/v1",  #  https://openrouter.ai/api/v1    |||||   https://api-inference.huggingface.co/v1/
-        api_key=chat_api_key2
+    history_text = ""
+    if chat_history:
+        history_text = "\n\n### Previous Conversation:\n"
+        for idx, (user_msg, bot_msg) in enumerate(chat_history, 1):
+            history_text += f"Turn {idx}:\nUser: {user_msg}\nTeacher: {bot_msg}\n"
+
+    identified_details_text = ""
+    if identified_details:
+        identified_details_text = "\n\n### Previously Identified Details:\n" + "\n".join(f"- {detail}" for detail in identified_details)
+
+    message_text = (
+        f"{history_text}{identified_details_text}\n\n"
+        f"Based on the image provided above, please evaluate the following description given by the child:\n"
+        f"'{user_details}'\n\n"
+        "You are a kind and encouraging teacher speaking to a child. Use simple, clear language. "
+        "Praise the child's correct observations and provide a gentle hint if something is missing. "
+        "Keep your feedback positive and easy to understand.\n\n"
+        "Focus on these evaluation criteria:\n"
+        "1. **Object Identification** – Did the child mention the main objects?\n"
+        "2. **Color & Shape Accuracy** – Were the colors and shapes described correctly?\n"
+        "3. **Clarity & Simplicity** – Was the description clear and easy to understand?\n"
+        "4. **Overall Communication** – How well did the child communicate their thoughts?\n\n"
+        "Note: As difficulty increases, the expected level of detail is higher. Evaluate accordingly.\n\n"
+        "Return your evaluation strictly as a JSON object with the following keys:\n"
+        "{\n"
+        "  \"scores\": {\n"
+        "    \"object_identification\": <number>,\n"
+        "    \"color_shape_accuracy\": <number>,\n"
+        "    \"clarity_simplicity\": <number>,\n"
+        "    \"overall_communication\": <number>\n"
+        "  },\n"
+        "  \"final_score\": <number>,\n"
+        "  \"feedback\": \"<string>\",\n"
+        "  \"hint\": \"<string>\",\n"
+        "  \"advance\": <boolean>\n"
+        "}\n\n"
+        "Do not include any additional text outside the JSON."
     )
 
-    stream = chat_client.chat.completions.create(
-        model="google/gemini-2.0-flash-thinking-exp:free",
-        messages=messages,
-        max_tokens=8192,
-        stream=True
-    )
+    # Remove the data:image/png;base64, prefix to get just the base64 string
+    base64_img = global_image_data_url.split(",")[1]
+
+    # Create a Gemini Vision Pro model
+    vision_model = GenerativeModel('gemini-2.0-flash-thinking-exp-01-21')
+
+    # Create the content with image and text using the correct parameters
+    # Use 'inline_data' instead of 'content' for the image part
+    image_part = Part(inline_data={"mime_type": "image/png", "data": base64.b64decode(base64_img)})
+    text_part = Part(text=message_text)
+    multimodal_content = Content(parts=[image_part, text_part])
 
-    response_text = ""
-    for chunk in stream:
-        response_text += chunk.choices[0].delta.content
-    return response_text
+    # Generate evaluation using the vision model
+    response = vision_model.generate_content(multimodal_content)
+
+    return response.text
+
+def evaluate_scores(evaluation_text, current_difficulty):
+    """
+    Parse the JSON evaluation and decide if the child advances.
+    The threshold scales with difficulty:
+       Very Simple: 70, Simple: 75, Moderate: 80, Detailed: 85, Very Detailed: 90.
+    """
+    try:
+        json_match = re.search(r'\{.*\}', evaluation_text, re.DOTALL)
+        if json_match:
+            json_str = json_match.group(0)
+            evaluation = json.loads(json_str)
+        else:
+            raise ValueError("No JSON object found in the response.")
+        final_score = evaluation.get("final_score", 0)
+        hint = evaluation.get("hint", "Keep trying!")
+        advance = evaluation.get("advance", False)
+        difficulty_thresholds = {
+            "Very Simple": 70,
+            "Simple": 75,
+            "Moderate": 80,
+            "Detailed": 85,
+            "Very Detailed": 90
+        }
+        current_threshold = difficulty_thresholds.get(current_difficulty, 70)
+        difficulty_mapping = {
+            "Very Simple": "Simple",
+            "Simple": "Moderate",
+            "Moderate": "Detailed",
+            "Detailed": "Very Detailed",
+            "Very Detailed": "Very Detailed"
+        }
+        if final_score >= current_threshold or advance:
+            new_difficulty = difficulty_mapping.get(current_difficulty, current_difficulty)
+            response_msg = (f"Great job! Your final score is {final_score}, which meets the target of {current_threshold}. "
+                            f"You've advanced to {new_difficulty} difficulty.")
+            return response_msg, new_difficulty
+        else:
+            response_msg = (f"Your final score is   {final_score} (\n target: {current_threshold}). {hint} \n "
+                            f"Please try again at the {current_difficulty} level.")
+            return response_msg, current_difficulty
+    except Exception as e:
+        return f"Error processing evaluation output: {str(e)}", current_difficulty
 
 def chat_respond(user_message, active_session, saved_sessions):
     """
-    Processes a new chat message. If no image has been generated yet, instructs the user to generate one.
-    Otherwise, sends the generated image and the user’s description to the vision language model for evaluation,
-    then appends the conversation to the active session's chat history.
+    Process a new chat message.
+    Evaluate the child's description. If the evaluation indicates advancement,
+    update the difficulty, generate a new image (resetting image and chat), and update the difficulty label.
     """
     if not active_session.get("image"):
         bot_message = "Please generate an image first."
+        updated_chat = active_session.get("chat", []) + [(user_message, bot_message)]
+        active_session["chat"] = updated_chat
+        return "", updated_chat, saved_sessions, active_session
+
+    chat_history = active_session.get("chat", [])
+    identified_details = active_session.get("identified_details", [])
+    raw_evaluation = compare_details_chat_fn(user_message, "", chat_history, identified_details)
+    current_difficulty = active_session.get("difficulty", "Very Simple")
+    evaluation_response, updated_difficulty = evaluate_scores(raw_evaluation, current_difficulty)
+    bot_message = evaluation_response
+
+    # If the child advanced, update difficulty and generate a new image
+    if updated_difficulty != current_difficulty:
+        # Update the active session's difficulty before generating a new prompt
+        active_session["difficulty"] = updated_difficulty
+        age = active_session.get("age", "3")
+        autism_level = active_session.get("autism_level", "Level 1")
+        topic_focus = active_session.get("topic_focus", "")
+        treatment_plan = active_session.get("treatment_plan", "")
+        new_image, new_active_session, new_sessions = generate_image_and_reset_chat(age, autism_level, topic_focus, treatment_plan, active_session, saved_sessions)
+        new_active_session["chat"].append(("System", f"You advanced to {updated_difficulty} difficulty! A new image has been generated for you."))
+        active_session = new_active_session
+        bot_message = f"You advanced to {updated_difficulty} difficulty! A new image has been generated for you."
+        saved_sessions = new_sessions
     else:
-        bot_message = compare_details_chat_fn(user_message)
+        updated_chat = active_session.get("chat", []) + [(user_message, bot_message)]
+        active_session["chat"] = updated_chat
 
-    updated_chat = active_session.get("chat", []) + [(user_message, bot_message)]
-    active_session["chat"] = updated_chat
-    return "", updated_chat, saved_sessions, active_session
+    return "", active_session["chat"], saved_sessions, active_session
 
 def update_sessions(saved_sessions, active_session):
     """
-    Combines finished sessions with the active session (if it exists)
-    so that the sidebar always displays the complete session details.
+    Combine finished sessions with the active session for display.
     """
     if active_session and active_session.get("prompt"):
         return saved_sessions + [active_session]
     return saved_sessions
 
 ##############################################
-# Dropdown Options for Difficulty and Level
-##############################################
-difficulty_options = ["Simple", "Average", "Detailed"]
-level_options = ["Level 1", "Level 2", "Level 3"]
-
-##############################################
-# Create the Gradio Interface (Single-Page) with a Sidebar for Session Details
+# Gradio Interface
 ##############################################
 with gr.Blocks() as demo:
-    active_session = gr.State({"prompt": None, "image": None, "chat": []})
+    # The active session now starts with difficulty "Very Simple"
+    active_session = gr.State({
+        "prompt": None,
+        "image": None,
+        "chat": [],
+        "treatment_plan": "",
+        "topic_focus": "",
+        "identified_details": [],
+        "difficulty": "Very Simple",
+        "age": "3",
+        "autism_level": "Level 1"
+    })
     saved_sessions = gr.State([])
 
     with gr.Column():
         gr.Markdown("# Image Generation & Chat Inference")
+        # Display current difficulty label
+        difficulty_label = gr.Markdown("**Current Difficulty:** Very Simple")
 
         # ----- Image Generation Section -----
         with gr.Column():
             gr.Markdown("## Generate Image")
-            gr.Markdown("Select options to create a custom prompt for image generation:")
+            gr.Markdown("Enter your age, select your autism level, specify a topic focus, and provide the treatment plan to generate an image based on the current difficulty level.")
             with gr.Row():
-                difficulty_dropdown = gr.Dropdown(label="Difficulty", choices=difficulty_options, value=difficulty_options[0])
-                # Changed age input from a dropdown to a text box
-                age_input = gr.Textbox(label="Age", placeholder="Enter your age...", value="3")
-                level_dropdown = gr.Dropdown(label="Level", choices=level_options, value=level_options[0])
-            # New textbox for extra details
-            extra_details_input = gr.Textbox(
-                label="Extra Details (optional)",
-                placeholder="Enter any additional details for the image...",
+                age_input = gr.Textbox(label="Age", placeholder="Enter age...", value="3")
+                autism_level_dropdown = gr.Dropdown(label="Autism Level", choices=["Level 1", "Level 2", "Level 3"], value="Level 1")
+
+            topic_focus_input = gr.Textbox(
+                label="Topic Focus",
+                placeholder="Enter a specific topic or detail to focus on (e.g., 'animals', 'emotions', 'daily routines')...",
+                lines=1
+            )
+
+            treatment_plan_input = gr.Textbox(
+                label="Treatment Plan",
+                placeholder="Enter the treatment plan to guide the image generation...",
                 lines=2
             )
             generate_btn = gr.Button("Generate Image")
             img_output = gr.Image(label="Generated Image")
             generate_btn.click(
                 generate_image_and_reset_chat,
-                inputs=[difficulty_dropdown, age_input, level_dropdown, extra_details_input, active_session, saved_sessions],
+                inputs=[age_input, autism_level_dropdown, topic_focus_input, treatment_plan_input, active_session, saved_sessions],
                 outputs=[img_output, active_session, saved_sessions]
             )
 
@@ -296,15 +316,12 @@ with gr.Blocks() as demo:
             gr.Markdown("## Chat about the Image")
             gr.Markdown(
                 "After generating an image, type details or descriptions about it. "
-                "Your message will be sent along with the image to a vision language model, "
-                "which will evaluate your description based on what it sees in the image. "
-                "The response will include a correctness percentage and hints if needed."
+                "Your message, along with the generated image and conversation history, will be sent for evaluation."
             )
             chatbot = gr.Chatbot(label="Chat History")
             with gr.Row():
                 chat_input = gr.Textbox(label="Your Message", placeholder="Type your description here...", show_label=False)
                 send_btn = gr.Button("Send")
-
             send_btn.click(
                 chat_respond,
                 inputs=[chat_input, active_session, saved_sessions],
@@ -322,11 +339,13 @@ with gr.Blocks() as demo:
         gr.Markdown(
             "This sidebar automatically saves finished chat sessions. "
             "Each session includes the prompt used, the generated image (as a data URL), "
-            "and the chat history (user messages and corresponding bot responses)."
+            "the topic focus, the treatment plan, the list of identified details, and the full chat history."
         )
         sessions_output = gr.JSON(label="Session Details", value={})
         active_session.change(update_sessions, inputs=[saved_sessions, active_session], outputs=sessions_output)
+        # Update the current difficulty label when active_session changes.
+        active_session.change(update_difficulty_label, inputs=[active_session], outputs=[difficulty_label])
         saved_sessions.change(update_sessions, inputs=[saved_sessions, active_session], outputs=sessions_output)
 
-# Launch the app.
+# Launch the app with public sharing enabled.
 demo.launch()