Update graph_tool.py

graph_tool.py

@@ -1,646 +1,126 @@
-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 base64
+from io import BytesIO
+import matplotlib.pyplot as plt
+import numpy as np
 import json
 
-@tool(return_direct=False)
-def Create_Graph_Tool(graph_config: str) -> str:
+def generate_plot(plot_config_json):
     """
-    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.
+    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.
     """
     try:
-        # 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
-        
+        # Parse the main JSON configuration
+        config = json.loads(plot_config_json)
     except json.JSONDecodeError as e:
-        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:
-        logger.error(f"Error in graph generation: {e}")
-        return f'<p style="color:red;">Error creating graph: {str(e)}</p>'
-
-# --- Tool Decision Engine ---
-class Tool_Decision_Engine:
-    """Uses LLM to intelligently decide when visualization tools would be beneficial"""
-    
-    def __init__(self, llm):
-        self.decision_llm = llm
-        self.decision_prompt = """Analyze this educational query and determine if creating a graph, chart, or visual representation would significantly enhance learning and understanding.
-
-Query: "{query}"
-
-Consider these factors:
-1. Would visualization make a concept clearer or easier to understand?
-2. Does the topic involve data, relationships, comparisons, or trends?
-3. Could a graph help illustrate abstract concepts concretely?
-4. For practice questions, would including visual elements be educational?
-
-Examples that BENEFIT from visualization:
-- Explaining mathematical functions or statistical concepts
-- Creating practice questions that involve data interpretation
-- Teaching about scientific trends or relationships
-- Comparing quantities, performance, or outcomes
-- Illustrating economic principles or business metrics
-
-Examples that do NOT need visualization:
-- Simple definitions or explanations
-- General conversation or greetings
-- Text-based study strategies
-- Qualitative discussions without data
-
-Answer with exactly: YES or NO
-
-Decision:"""
-
-    def should_use_visualization(self, query: str) -> bool:
-        """Use LLM reasoning to determine if visualization would be beneficial"""
-        try:
-            # Create decision prompt
-            decision_query = self.decision_prompt.format(query=query)
-            
-            # Get LLM decision
-            decision_response = self.decision_llm.invoke(decision_query)
-            
-            # Parse response - look for YES/NO
-            decision_text = decision_response.strip().upper()
-            
-            # Log the decision for debugging
-            logger.info(f"Tool decision for '{query[:50]}...': {decision_text}")
-            
-            return "YES" in decision_text and "NO" not in decision_text
-            
-        except Exception as e:
-            logger.error(f"Error in tool decision making: {e}")
-            # Default to no tools if decision fails
-            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 LLM-based tool decision making"""
-    
-    def __init__(self):
-        self.llm = Qwen25SmallLLM(model_path="Qwen/Qwen2.5-3B-Instruct")
-        self.tool_decision_engine = Tool_Decision_Engine(self.llm)
-        self.memory = ConversationBufferWindowMemory(
-            memory_key="chat_history",
-            return_messages=True,
-            k=5
-        )
-        
-    def should_use_tools(self, query: str) -> bool:
-        """Use LLM reasoning to determine if tools are needed"""
-        return self.tool_decision_engine.should_use_visualization(query)
+        return f'<p style="color:red;">Error parsing JSON configuration: {e}</p>'
     
-    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:
-        # 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}")
-        
-        # Generate response with modern agent
-        response = generate_response_with_agent(message)
-        
-        # Log final metrics
-        try:
-            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...")
+    # Extract parameters with defaults
     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()
-        
+        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", "")
     except Exception as e:
-        logger.error(f"Agent warmup failed: {e}")
+        return f'<p style="color:red;">Error extracting configuration parameters: {e}</p>'
     
-# --- UI: Interface Creation ---
-def create_interface():
-    """Creates and configures the complete Gradio interface."""
+    # 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>'
     
-    # 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}")
+    if not isinstance(labels, list):
+        return '<p style="color:red;">Labels must be a list.</p>'
     
-    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>')
-            
-            # 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"
-                )
-            
-            # 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")
-            
-            # 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()
+        fig, ax = plt.subplots(figsize=(10, 6))
         
-        interface = create_interface()
-        interface.queue()
-        interface.launch(
-            server_name="0.0.0.0",
-            share=True,
-            debug=True,
-            favicon_path="assets/favicon.ico"
-        )
+        # Extract keys and values from the data dictionary
+        x_data = list(data.keys())
+        y_data = list(data.values())
         
+        # Ensure y_data contains numeric values
+        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)
+            
+            # 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)
+            
+        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
+            
+        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>
+        '''
+         
     except Exception as e:
-        logger.error(f"❌ Failed to launch Mimir: {e}")
-        raise
+        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))