Update graph_tool.py

graph_tool.py

@@ -1,126 +1,621 @@
-import base64
-from io import BytesIO
-import matplotlib.pyplot as plt
-import numpy as np
+import gradio as gr
+from graph_tool import generate_plot
+from metrics import MimirMetrics
+import os
+
+os.environ['HF_HOME'] = '/tmp/huggingface'
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/huggingface'
+os.environ['HF_DATASETS_CACHE'] = '/tmp/huggingface'
+
+import time
+from dotenv import load_dotenv
+import logging
+import re
+from langchain_core.tools import tool
+from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnableBranch
+from langgraph.prebuilt import create_react_agent
+from langchain.memory import ConversationBufferWindowMemory
+from typing import Optional, List, Any, Type
+from pydantic import BaseModel, Field
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import torch
+
+load_dotenv(".env")
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
+print("Environment variables loaded.")
+
+# --- Environment and Logging Setup ---
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Support both token names for flexibility
+hf_token = HF_TOKEN
+if not hf_token:
+    logger.warning("Neither HF_TOKEN nor HUGGINGFACEHUB_API_TOKEN is set, the application may not work.")
+
+metrics_tracker = MimirMetrics(save_file="Mimir_metrics.json")
+
 import json
 
-def generate_plot(plot_config_json):
+@tool(return_direct=False)
+def Create_Graph_Tool(graph_config: str) -> str:
     """
-    Generates a plot (bar, line, or pie) and returns it as an HTML-formatted Base64-encoded image string.
-    
-    Args:
-        plot_config_json (str): JSON string containing all plot configuration:
-            {
-                "data": {"key1": value1, "key2": value2, ...},
-                "labels": ["label1", "label2", ...],
-                "plot_type": "bar|line|pie",
-                "title": "Plot Title",
-                "x_label": "X Axis Label" (optional),
-                "y_label": "Y Axis Label" (optional)
-            }
-        
-    Returns:
-        str: An HTML img tag with Base64-encoded plot image.
+    Creates educational graphs and charts to help explain concepts to students.
+    
+    Use this tool ONLY when teaching concepts that would benefit from visual representation, such as:
+    - Mathematical functions and relationships (quadratic equations, exponential growth)
+    - Statistical distributions and data analysis (normal curves, survey results)
+    - Scientific trends and comparisons (temperature changes, population growth)
+    - Economic models and business metrics (profit over time, market shares)
+    - Grade distributions or performance analysis (test score ranges)
+    - Any quantitative concept that's clearer with visualization
+    
+    Input should be a JSON string with this structure:
+    {
+        "data": {"Category A": 25, "Category B": 40, "Category C": 35},
+        "plot_type": "bar",
+        "title": "Student Performance by Subject", 
+        "x_label": "Subjects",
+        "y_label": "Average Score",
+        "educational_context": "This visualization helps students see performance patterns across subjects"
+    }
+    
+    Plot types:
+    - "bar": Best for comparing categories, showing distributions, or discrete data
+    - "line": Best for showing trends over time or continuous relationships  
+    - "pie": Best for showing parts of a whole or proportions
+    
+    Always create meaningful educational data that illustrates the concept you're teaching.
+    Include educational_context to explain why the visualization helps learning.
     """
     try:
-        # Parse the main JSON configuration
-        config = json.loads(plot_config_json)
+        # Validate it's proper JSON
+        config = json.loads(graph_config)
+        
+        # Add educational context if provided
+        educational_context = config.get("educational_context", "")
+        
+        # Call your generate_plot function
+        graph_html = generate_plot(graph_config)
+        
+        # Add educational context if provided
+        if educational_context:
+            context_html = f'<div style="margin: 10px 0; padding: 10px; background: #f8f9fa; border-left: 4px solid #007bff; font-style: italic;">💡 {educational_context}</div>'
+            return context_html + graph_html
+        
+        return graph_html
+        
     except json.JSONDecodeError as e:
-        return f'<p style="color:red;">Error parsing JSON configuration: {e}</p>'
-    
-    # Extract parameters with defaults
-    try:
-        data = config.get("data", {})
-        labels = config.get("labels", [])
-        plot_type = config.get("plot_type", "bar")
-        title = config.get("title", "Untitled Plot")
-        x_label = config.get("x_label", "")
-        y_label = config.get("y_label", "")
+        logger.error(f"Invalid JSON provided to graph tool: {e}")
+        return '<p style="color:red;">Graph generation failed - invalid JSON format</p>'
     except Exception as e:
-        return f'<p style="color:red;">Error extracting configuration parameters: {e}</p>'
+        logger.error(f"Error in graph generation: {e}")
+        return f'<p style="color:red;">Error creating graph: {str(e)}</p>'
+
+# --- Tool Control Logic ---
+class Tool_Manager:
+    """Manages when tools should be available"""
     
-    # Validate inputs
-    if not isinstance(data, dict):
-        return '<p style="color:red;">Data must be a dictionary with keys as labels and values as numbers.</p>'
+    def __init__(self):
+        self.graph_keywords = [
+            'graph', 'chart', 'plot', 'visualiz', 'diagram', 'show me',
+            'data', 'statistic', 'trend', 'distribution', 'comparison'
+        ]
     
-    if not isinstance(labels, list):
-        return '<p style="color:red;">Labels must be a list.</p>'
+    def needs_visualization(self, query: str) -> bool:
+        """Determine if query requires visualization tools"""
+        query_lower = query.lower()
+        
+        # Check for explicit visualization requests
+        visualization_patterns = [
+            r'\b(graph|chart|plot|visualiz)\w*',
+            r'\bshow\s+me\s+(a\s+)?(graph|chart)',
+            r'\b(data|statistic)\w*\s+(analys|visual)',
+            r'\bcreate\s+(a\s+)?(graph|chart|diagram)',
+            r'\bmake\s+(a\s+)?(graph|chart)',
+            r'\bdraw\s+(a\s+)?(graph|chart)'
+        ]
+        
+        for pattern in visualization_patterns:
+            if re.search(pattern, query_lower):
+                return True
+        
+        return False
+
+# --- System Prompt ---
+SYSTEM_PROMPT = """You are Mimir, an expert multi-concept tutor designed to facilitate genuine learning and understanding. Your primary mission is to guide students through the learning process rather than providing direct answers to academic work.
+
+## Core Educational Principles
+- Provide comprehensive, educational responses that help students truly understand concepts
+- Use minimal formatting, with markdown bolding reserved for **key terms** only
+- Prioritize teaching methodology over answer delivery
+- Foster critical thinking and independent problem-solving skills
+
+## Tone and Communication Style
+- Maintain an engaging, friendly tone appropriate for high school students
+- Write at a reading level that is accessible yet intellectually stimulating
+- Be supportive and encouraging without being condescending
+- Never use crude language or content inappropriate for an educational setting
+- Avoid preachy, judgmental, or accusatory language
+- Skip flattery and respond directly to questions
+- Do not use emojis or actions in asterisks unless specifically requested
+- Present critiques and corrections kindly as educational opportunities
+
+## Academic Integrity Approach
+You recognize that students may seek direct answers to homework, assignments, or test questions. Rather than providing complete solutions or making accusations about intent, you should:
+
+- **Guide through processes**: Break down problems into conceptual components and teach underlying principles
+- **Ask clarifying questions**: Understand what the student already knows and where their confusion lies
+- **Provide similar examples**: Work through analogous problems that demonstrate the same concepts without directly solving their specific assignment
+- **Encourage original thinking**: Help students develop their own reasoning and analytical skills
+- **Suggest study strategies**: Recommend effective learning approaches for the subject matter
+
+## Visual Learning Enhancement
+You have the ability to create graphs and charts to enhance your explanations. Use this capability proactively when:
+- Explaining mathematical concepts (functions, distributions, relationships)
+- Teaching statistical analysis or data interpretation
+- Discussing scientific trends, patterns, or experimental results
+- Comparing different options, outcomes, or scenarios
+- Illustrating economic principles, business metrics, or financial concepts
+- Showing survey results, demographic data, or research findings
+- Demonstrating any concept where visualization aids comprehension
+
+**Important**: Only use the graph tool when visualization would genuinely help explain a concept. For general conversation, explanations, or questions that don't involve data or relationships, respond normally without tools.
+
+## Response Guidelines
+- **For math problems**: Explain concepts, provide formula derivations, and guide through problem-solving steps without computing final numerical answers
+- **For multiple-choice questions**: Discuss the concepts being tested and help students understand how to analyze options rather than identifying the correct choice
+- **For essays or written work**: Discuss research strategies, organizational techniques, and critical thinking approaches rather than providing content or thesis statements
+- **For factual questions**: Provide educational context and encourage students to synthesize information rather than stating direct answers
+
+## Communication Guidelines
+- Maintain a supportive, non-judgmental tone in all interactions
+- Assume positive intent while redirecting toward genuine learning
+- Use Socratic questioning to promote discovery and critical thinking
+- Celebrate understanding and progress in the learning process
+- Encourage students to explain their thinking and reasoning
+- Provide honest, accurate feedback even when it may not be what the student wants to hear
+
+Your goal is to be an educational partner who empowers students to succeed through understanding, not a service that completes their work for them."""
+
+# --- LLM Class Unchanged ---
+logger = logging.getLogger(__name__)
+
+class Qwen25SmallLLM:
+    def __init__(self, model_path: str = "Qwen/Qwen2.5-3B-Instruct", use_4bit: bool = True):
+        logger.info(f"Loading model: {model_path} (use_4bit={use_4bit})")
+
+        try:
+            # Load tokenizer
+            self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+            if use_4bit:
+                quant_config = BitsAndBytesConfig(
+                    load_in_4bit=True,
+                    bnb_4bit_compute_dtype=torch.bfloat16,
+                    bnb_4bit_use_double_quant=True,
+                    bnb_4bit_quant_type="nf4"
+                )
+                logger.info("Using 4-bit quantization with BitsAndBytes")
+            else:
+                quant_config = BitsAndBytesConfig(
+                    load_in_8bit=True,
+                    llm_int8_enable_fp32_cpu_offload=True
+                )
+                logger.info("Using 8-bit quantization with BitsAndBytes")
+
+            # Try quantized load
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_path,
+                quantization_config=quant_config,
+                device_map="auto",
+                torch_dtype=torch.bfloat16,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True
+            )
+
+        except Exception as e:
+            logger.warning(f"Quantized load failed, falling back: {e}")
+            self._load_fallback_model(model_path)
+
+        # Ensure pad token
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+    def _load_fallback_model(self, model_path: str):
+        """Fallback if quantization fails."""
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto" if torch.cuda.is_available() else None,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True
+        )
+
+    def invoke(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        try:
+            messages = [
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": prompt}
+            ]
+            text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+            inputs = self.tokenizer([text], return_tensors="pt", padding=True, truncation=True, max_length=2048)
+            if torch.cuda.is_available():
+                inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
+
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_new_tokens=800,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_p=0.9,
+                    top_k=50,
+                    repetition_penalty=1.1,
+                    pad_token_id=self.tokenizer.eos_token_id
+                )
+
+            new_tokens = [out[len(inp):] for inp, out in zip(inputs.input_ids, outputs)]
+            return self.tokenizer.batch_decode(new_tokens, skip_special_tokens=True)[0].strip()
+
+        except Exception as e:
+            logger.error(f"Generation error: {e}")
+            return f"[Error generating response: {str(e)}]"
+
+# --- Modern Agent Implementation ---
+class Educational_Agent:
+    """Modern LangChain agent with proper tool control"""
     
+    def __init__(self):
+        self.llm = Qwen25SmallLLM(model_path="Qwen/Qwen2.5-3B-Instruct")
+        self.tool_manager = Tool_Manager()
+        self.memory = ConversationBufferWindowMemory(
+            memory_key="chat_history",
+            return_messages=True,
+            k=5
+        )
+        
+    def should_use_tools(self, query: str) -> bool:
+        """Determine if this query needs visualization tools"""
+        return self.tool_manager.needs_visualization(query)
+    
+    def create_prompt_template(self, has_tools: bool = False):
+        """Create prompt template based on whether tools are available"""
+        if has_tools:
+            system_content = SYSTEM_PROMPT + "\n\nYou have access to graph creation tools. Use them when visualization would help explain concepts."
+        else:
+            system_content = SYSTEM_PROMPT + "\n\nRespond using your knowledge without any tools."
+            
+        return ChatPromptTemplate.from_messages([
+            ("system", system_content),
+            ("human", "{input}")
+        ])
+    
+    def process_with_tools(self, query: str) -> str:
+        """Process query with tools available"""
+        try:
+            # Create agent with tools
+            tools = [Create_Graph_Tool]
+            
+            # Use create_react_agent for better control
+            agent = create_react_agent(
+                self.llm,
+                tools,
+                state_modifier=self.create_prompt_template(has_tools=True)
+            )
+            
+            response = agent.invoke({"messages": [HumanMessage(content=query)]})
+            
+            # Extract the final message content
+            if response and "messages" in response:
+                final_message = response["messages"][-1]
+                if hasattr(final_message, 'content'):
+                    return final_message.content
+                else:
+                    return str(final_message)
+            
+            return str(response)
+            
+        except Exception as e:
+            logger.error(f"Error in tool processing: {e}")
+            return f"I apologize, but I encountered an error while processing your request: {str(e)}"
+    
+    def process_without_tools(self, query: str) -> str:
+        """Process query without tools"""
+        try:
+            response = self.llm.invoke(query)
+            return response
+        except Exception as e:
+            logger.error(f"Error in normal processing: {e}")
+            return f"I apologize, but I encountered an error: {str(e)}"
+    
+    def chat(self, message: str) -> str:
+        """Main chat interface with conditional tool usage"""
+        try:
+            # Determine if tools are needed
+            if self.should_use_tools(message):
+                logger.info("Query requires visualization - enabling tools")
+                return self.process_with_tools(message)
+            else:
+                logger.info("Query doesn't need tools - responding normally")
+                return self.process_without_tools(message)
+                
+        except Exception as e:
+            logger.error(f"Error in chat processing: {e}")
+            return f"I apologize, but I encountered an error: {str(e)}"
+
+# --- Global Agent Instance ---
+agent = None
+
+def get_agent():
+    """Get or create the educational agent."""
+    global agent
+    if agent is None:
+        agent = Educational_Agent()
+    return agent
+
+# --- UI: MathJax Configuration ---
+mathjax_config = '''
+<script>
+window.MathJax = {
+  tex: {
+    inlineMath: [['\\\\(', '\\\\)']],
+    displayMath: [['$', '$'], ['\\\\[', '\\\\]']],
+    packages: {'[+]': ['ams']}
+  },
+  svg: {fontCache: 'global'},
+  startup: {
+    ready: () => {
+      MathJax.startup.defaultReady();
+      // Re-render math when new content is added
+      const observer = new MutationObserver(function(mutations) {
+        MathJax.typesetPromise();
+      });
+      observer.observe(document.body, {childList: true, subtree: true});
+    }
+  }
+};
+</script>
+'''
+
+# --- HTML Head Content ---
+html_head_content = '''
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>Mimir - AI Educational Assistant</title>
+'''
+
+# --- Force Light Mode Script ---
+force_light_mode = '''
+<script>
+// Force light theme in Gradio
+window.addEventListener('DOMContentLoaded', function () {
+    const gradioURL = window.location.href;
+    const url = new URL(gradioURL);
+    const currentTheme = url.searchParams.get('__theme');
+    
+    if (currentTheme !== 'light') {
+        url.searchParams.set('__theme', 'light');
+        window.location.replace(url.toString());
+    }
+});
+</script>
+'''
+
+# --- Core Logic Functions ---
+def smart_truncate(text, max_length=3000):
+    """Truncates text intelligently to the last full sentence or word."""
+    if len(text) <= max_length:
+        return text
+    
+    # Try to split by sentence
+    sentences = re.split(r'(?<=[.!?])\s+', text[:max_length])
+    if len(sentences) > 1:
+        return ' '.join(sentences[:-1]) + "... [Response truncated - ask for continuation]"
+    # Otherwise, split by word
+    words = text[:max_length].split()
+    return ' '.join(words[:-1]) + "... [Response truncated]"
+
+def generate_response_with_agent(message, max_retries=3):
+    """Generate response using modern agent with proper tool control."""
+    
+    for attempt in range(max_retries):
+        try:
+            # Get the agent
+            current_agent = get_agent()
+            
+            # Use the agent's chat method with conditional tool usage
+            response = current_agent.chat(message)
+            
+            return smart_truncate(response)
+            
+        except Exception as e:
+            logger.error(f"Agent error (attempt {attempt + 1}): {e}")
+            if attempt < max_retries - 1:
+                time.sleep(2)
+                continue
+            else:
+                return f"I apologize, but I encountered an error while processing your message: {str(e)}"
+
+def chat_response(message, history=None):
+    """Process chat message and return response."""
     try:
-        fig, ax = plt.subplots(figsize=(10, 6))
+        # Track metrics with timing context
+        start_time = time.time()
+        timing_context = {
+            'start_time': start_time,
+            'chunk_count': 0,
+            'provider_latency': 0.0
+        }
+        
+        try:
+            # Log start of interaction
+            metrics_tracker.log_interaction(
+                query=message,
+                response="", 
+                timing_context=timing_context,
+                error_occurred=False
+            )
+            logger.info("Metrics interaction logged successfully")
+        except Exception as metrics_error:
+            logger.error(f"Error in metrics_tracker.log_interaction: {metrics_error}")
         
-        # Extract keys and values from the data dictionary
-        x_data = list(data.keys())
-        y_data = list(data.values())
+        # Generate response with modern agent
+        response = generate_response_with_agent(message)
         
-        # Ensure y_data contains numeric values
+        # Log final metrics
         try:
-            y_data = [float(val) for val in y_data]
-        except (ValueError, TypeError):
-            return '<p style="color:red;">All data values must be numeric.</p>'
-        
-        if plot_type == 'bar':
-            bars = ax.bar(x_data, y_data)
-            ax.set_xlabel(x_label)
-            ax.set_ylabel(y_label)
+            metrics_tracker.log_interaction(
+                query=message,
+                response=response,
+                timing_context=timing_context,
+                error_occurred=False
+            )
+        except Exception as metrics_error:
+            logger.error(f"Error in final metrics logging: {metrics_error}")
+        
+        return response
+        
+    except Exception as e:
+        logger.error(f"Error in chat_response: {e}")
+        return f"I apologize, but I encountered an error while processing your message: {str(e)}"
+
+def respond_and_update(message, history):
+    """Main function to handle user submission."""
+    if not message.strip():
+        return history, ""
+    
+    # Add user message to history
+    history.append({"role": "user", "content": message})
+    yield history, ""
+
+    # Generate response
+    response = chat_response(message)
+    
+    history.append({"role": "assistant", "content": response})
+    yield history, ""
+
+def clear_chat():
+    """Clear the chat history."""
+    global agent
+    if agent is not None:
+        agent.memory.clear()
+    return [], ""
+
+def warmup_agent():
+    """Warm up the agent with a test query to preload everything."""
+    logger.info("Warming up agent with test query...")
+    try:
+        current_agent = get_agent()
+        
+        # Run a simple test query
+        test_response = current_agent.chat("Hello, this is a warmup test.")
+        logger.info(f"Agent warmup completed successfully! Test response length: {len(test_response)} chars")
+        
+        # Clear any test data from memory
+        current_agent.memory.clear()
+        
+    except Exception as e:
+        logger.error(f"Agent warmup failed: {e}")
+    
+# --- UI: Interface Creation ---
+def create_interface():
+    """Creates and configures the complete Gradio interface."""
+    
+    # Read CSS file
+    custom_css = ""
+    try:
+        with open("styles.css", "r", encoding="utf-8") as css_file:
+            custom_css = css_file.read()
+    except FileNotFoundError:
+        logger.warning("styles.css file not found, using default styling")
+    except Exception as e:
+        logger.warning(f"Error reading styles.css: {e}")
+    
+    with gr.Blocks(
+        title="Mimir", 
+        fill_width=True, 
+        fill_height=True,
+        theme=gr.themes.Origin()
+    ) as demo:
+        # Add head content and MathJax
+        gr.HTML(html_head_content)
+        gr.HTML(force_light_mode)
+        gr.HTML('<script id="MathJax-script" async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>')
+        gr.HTML(mathjax_config)
+        
+        with gr.Column(elem_classes=["main-container"]):
+            # Title Section
+            gr.HTML('<div class="title-header"><h1> Mimir 🎓</h1></div>')
             
-            # Add value labels on top of bars
-            for bar, value in zip(bars, y_data):
-                height = bar.get_height()
-                ax.text(bar.get_x() + bar.get_width()/2., height,
-                       f'{value}', ha='center', va='bottom')
-                       
-        elif plot_type == 'line':
-            ax.plot(x_data, y_data, marker='o', linewidth=2, markersize=6)
-            ax.set_xlabel(x_label)
-            ax.set_ylabel(y_label)
-            ax.grid(True, alpha=0.3)
+            # Chat Section
+            with gr.Row():
+                chatbot = gr.Chatbot(
+                    type="messages",
+                    show_copy_button=True,
+                    show_share_button=False,
+                    avatar_images=None,
+                    elem_id="main-chatbot",
+                    container=False,
+                    scale=1,
+                    height="70vh"
+                )
             
-        elif plot_type == 'pie':
-            # For pie charts, use labels parameter if provided, otherwise use data keys
-            pie_labels = labels if len(labels) == len(y_data) else x_data
-            wedges, texts, autotexts = ax.pie(y_data, labels=pie_labels, autopct='%1.1f%%', 
-                                            startangle=90, textprops={'fontsize': 10})
-            ax.axis('equal')  # Equal aspect ratio ensures that pie is drawn as a circle
-            # Don't set x/y labels for pie charts as they don't make sense
+            # Input Section
+            with gr.Row(elem_classes=["input-controls"]):
+                msg = gr.Textbox(
+                    placeholder="Ask me about math, research, study strategies, or any educational topic...",
+                    show_label=False,
+                    lines=6,
+                    max_lines=8,
+                    elem_classes=["input-textbox"],
+                    container=False,
+                    scale=4
+                )
+                with gr.Column(elem_classes=["button-column"], scale=1):
+                    send = gr.Button("Send", elem_classes=["send-button"], size="sm")
+                    clear = gr.Button("Clear", elem_classes=["clear-button"], size="sm")
             
-        else:
-            return f'<p style="color:red;">Invalid plot_type: {plot_type}. Choose "bar", "line", or "pie".</p>'
-        
-        ax.set_title(title, fontsize=14, fontweight='bold', pad=20)
-        
-        # Improve layout
-        plt.tight_layout()
-        
-        # Save plot to a BytesIO buffer in memory
-        buf = BytesIO()
-        plt.savefig(buf, format='png', bbox_inches='tight', dpi=150, 
-                   facecolor='white', edgecolor='none')
-        plt.close(fig)  # Close the plot to free up memory
-        
-        # Encode the image data to a Base64 string
-        img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
-        
-        # Return HTML img tag with proper styling
-        return f'''
-        <div style="text-align: center; margin: 20px 0;">
-            <img src="data:image/png;base64,{img_base64}" 
-                 style="max-width: 100%; height: auto; border-radius: 8px; box-shadow: 0 2px 10px rgba(0,0,0,0.1);" 
-                 alt="{title}" />
-        </div>
-        '''
-         
+            # Event handlers
+            msg.submit(respond_and_update, [msg, chatbot], [chatbot, msg])
+            send.click(respond_and_update, [msg, chatbot], [chatbot, msg])
+            clear.click(clear_chat, outputs=[chatbot, msg])
+            
+            # Apply CSS at the very end
+            gr.HTML(f'<style>{custom_css}</style>')
+        
+        return demo
+
+# --- Main Execution ---
+if __name__ == "__main__":
+    try:
+        logger.info("=" * 50)
+        logger.info("Starting Mimir Application")
+        logger.info("=" * 50)
+        
+        # Step 1: Preload the model and agent
+        logger.info("Loading AI model...")
+        start_time = time.time()
+        agent = Educational_Agent()
+        load_time = time.time() - start_time
+        logger.info(f"Model loaded successfully in {load_time:.2f} seconds")
+        
+        # Step 2: Warm up the model
+        logger.info("Warming up model...")
+        warmup_agent()
+        
+        interface = create_interface()
+        interface.queue()
+        interface.launch(
+            server_name="0.0.0.0",
+            share=True,
+            debug=True,
+            favicon_path="assets/favicon.ico"
+        )
+        
     except Exception as e:
-        plt.close('all')  # Clean up any open figures
-        return f'<p style="color:red;">Error generating plot: {str(e)}</p>'
-
-# Example usage:
-# plot_config = {
-#     "data": {"A": 10, "B": 20, "C": 15},
-#     "labels": ["Category A", "Category B", "Category C"],
-#     "plot_type": "bar",
-#     "title": "Sample Bar Chart",
-#     "x_label": "Categories",
-#     "y_label": "Values"
-# }
-# result = generate_plot_single_input(json.dumps(plot_config))
+        logger.error(f"❌ Failed to launch Mimir: {e}")
+        raise