Update app.py

app.py

@@ -11,11 +11,12 @@ import time
 from dotenv import load_dotenv
 import logging
 import re
-from langchain.tools import BaseTool
-from langchain.agents import initialize_agent, AgentType
+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 langchain.schema import SystemMessage
-from langchain.llms.base import LLM
 from typing import Optional, List, Any, Type
 from pydantic import BaseModel, Field
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
@@ -36,103 +37,117 @@ if not hf_token:
 
 metrics_tracker = MimirMetrics(save_file="Mimir_metrics.json")
 
-from langchain.tools import Tool
 import json
 
-def create_educational_graph_tool():
-    """
-    Creates a graph tool for the AI to autonomously generate educational visualizations.
+@tool(return_direct=False)
+def Create_Graph_Tool(graph_config: str) -> str:
     """
+    Creates educational graphs and charts to help explain concepts to students.
     
-    def graph_wrapper(graph_config: str) -> str:
-        """
-        Wrapper that calls your generate_plot function with the JSON config.
-        
-        Args:
-            graph_config (str): JSON string with plot configuration
-        
-        Returns:
-            str: HTML with embedded graph
-        """
-        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 (which now takes single JSON input)
-            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:
-            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>'
+    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
     
-    return Tool(
-        name="create_graph",
-        description="""Create educational graphs and charts to help explain concepts to students.
+    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:
+        # Validate it's proper JSON
+        config = json.loads(graph_config)
         
-        Use this tool 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
+        # Add educational context if provided
+        educational_context = config.get("educational_context", "")
         
-        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"
-        }
+        # Call your generate_plot function
+        graph_html = generate_plot(graph_config)
         
-        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
+        # 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
         
-        Always create meaningful educational data that illustrates the concept you're teaching.
-        Include educational_context to explain why the visualization helps learning.
+        return graph_html
         
-        Examples of when to use:
-        - Student asks about probability → Create histogram showing dice roll outcomes
-        - Teaching compound interest → Line chart showing money growth over time
-        - Explaining survey methods → Pie chart of sample demographic breakdown
-        - Discussing functions → Line plot comparing linear vs exponential growth
-        - Analyzing test performance → Bar chart of score distributions
-        """,
-        func=graph_wrapper
-    )
-
-# If you want to keep using your existing CreateGraphTool approach, 
-# update it to work with the new single-input generate_plot function:
-class CreateGraphToolFixed(BaseTool):
-    name: str = "create_graph"
-    description: str = """Generate educational graphs to supplement explanations. Create realistic data that illustrates educational concepts."""
+    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 _run(self, tool_input: str) -> str:
-        """Single input version that works with your updated generate_plot function."""
+    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:
-            # tool_input should be the JSON config string
-            return generate_plot(tool_input)
+            # 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 graph tool: {e}")
-            return f'<p style="color:red;">Error creating graph: {str(e)}</p>'
-
+            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.
@@ -162,12 +177,8 @@ You recognize that students may seek direct answers to homework, assignments, or
 - **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
-
-Thought: You should always think about what to do; do not use any tool if it is not needed.
-
+## 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
@@ -176,42 +187,13 @@ You have the ability to create graphs and charts to enhance your explanations. U
 - Showing survey results, demographic data, or research findings
 - Demonstrating any concept where visualization aids comprehension
 
-**When to create graphs:**
-- The concept involves numerical data or relationships
-- Visual representation would clarify a complex idea
-- Students benefit from seeing patterns or comparisons
-- You're teaching about graph interpretation itself
-- The topic involves trends, distributions, or proportions
-- A multiple choice question requires a visual, such as a graph, to be answered.
-
-**How to create graphs:**
-Generate realistic, educational data that illustrates your teaching point. Create meaningful examples that help students understand the underlying concepts, not just random numbers.
-
-Example: When explaining normal distribution, create a graph showing test scores distributed normally around a mean, with appropriate labels and educational context.
-
-
-## Tool Usage
-You have access to a create_graph tool. Use this tool naturally when a visual representation would enhance understanding or when discussing concepts that involve data, relationships, patterns, or quantitative information. Consider creating graphs for:
-- Mathematical concepts (functions, distributions, relationships)
-- Statistical examples and explanations
-- Scientific data and relationships
-- Practice problems involving graph interpretation
-- Comparative analyses
-- Economic models or business concepts
-- Any situation where visualization aids comprehension
-
-When using the create_graph tool, format data as JSON strings:
-- data_json: '{"Category1": 25, "Category2": 40, "Category3": 35}'
-- labels_json: '["Category1", "Category2", "Category3"]'
-
-Thought: You should always think about what to do; do not use any tool if it is not needed.
+**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
-- Use graphs naturally when they would clarify or enhance your explanations
 
 ## Communication Guidelines
 - Maintain a supportive, non-judgmental tone in all interactions
@@ -223,26 +205,11 @@ Thought: You should always think about what to do; do not use any tool if it is
 
 Your goal is to be an educational partner who empowers students to succeed through understanding, not a service that completes their work for them."""
 
-# --- Improved LangChain Setup ---
-# Global flag to track system prompt initialization
-system_prompt_initialized = False
-
-def initialize_system_prompt(agent):
-    """Initialize the system prompt as a SystemMessage in memory."""
-    global system_prompt_initialized
-    if not system_prompt_initialized:
-        system_message = SystemMessage(content=SYSTEM_PROMPT)
-        agent.memory.chat_memory.add_message(system_message)
-        system_prompt_initialized = True
-
+# --- LLM Class Unchanged ---
 logger = logging.getLogger(__name__)
 
-class Qwen25SmallLLM(LLM):
-    model: Any = None
-    tokenizer: Any = None
-
+class Qwen25SmallLLM:
     def __init__(self, model_path: str = "Qwen/Qwen2.5-3B-Instruct", use_4bit: bool = True):
-        super().__init__()
         logger.info(f"Loading model: {model_path} (use_4bit={use_4bit})")
 
         try:
@@ -292,7 +259,7 @@ class Qwen25SmallLLM(LLM):
             low_cpu_mem_usage=True
         )
 
-    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+    def invoke(self, prompt: str, stop: Optional[List[str]] = None) -> str:
         try:
             messages = [
                 {"role": "system", "content": SYSTEM_PROMPT},
@@ -323,94 +290,98 @@ class Qwen25SmallLLM(LLM):
             logger.error(f"Generation error: {e}")
             return f"[Error generating response: {str(e)}]"
 
-    @property
-    def _llm_type(self) -> str:
-        return "qwen25_small"
-
-
-# Example of how the AI should use the tool
-def example_usage_for_ai():
-    """
-    This shows how the AI should autonomously create graphs in its responses.
-    The AI doesn't wait for user data - it creates meaningful educational examples.
-    """
+# --- Modern Agent Implementation ---
+class Educational_Agent:
+    """Modern LangChain agent with LLM-based tool decision making"""
     
-    # Example: Teaching about normal distribution
-    example_config = {
-        "data": {
-            "Below 60": 5,
-            "60-70": 15, 
-            "70-80": 25,
-            "80-90": 35,
-            "90-100": 20
-        },
-        "plot_type": "bar",
-        "title": "Typical Test Score Distribution",
-        "x_label": "Score Range",
-        "y_label": "Number of Students",
-        "educational_context": "This shows how test scores often follow a bell-curve pattern, with most students scoring in the middle range."
-    }
+    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)
     
-    # Example: Teaching about compound interest
-    compound_interest_example = {
-        "data": {
-            "Year 1": 1000,
-            "Year 5": 1276,
-            "Year 10": 1629,
-            "Year 15": 2079,
-            "Year 20": 2653
-        },
-        "plot_type": "line", 
-        "title": "Compound Interest Growth ($1000 at 5% Annual)",
-        "x_label": "Time (Years)",
-        "y_label": "Account Value ($)",
-        "educational_context": "Notice how the growth accelerates over time - this is the power of compound interest!"
-    }
+    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}")
+        ])
     
-    return "These examples show how the AI creates educational data to illustrate concepts"
+    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 LangChain agent."""
+    """Get or create the educational agent."""
     global agent
     if agent is None:
-        agent = create_langchain_agent()
+        agent = Educational_Agent()
     return agent
 
-def create_langchain_agent():
-    """Factory to build the LangChain agent with memory and tools."""
-    try:
-        # Initialize your LLM
-        llm = Qwen25SmallLLM(model_path="Qwen/Qwen2.5-3B-Instruct")
-
-        # Memory
-        memory = ConversationBufferWindowMemory(
-            memory_key="chat_history", 
-            return_messages=True,
-            k=5  # keep last 5 exchanges
-        )
-
-        # Tools (graph tool, etc.)
-        tools = [create_educational_graph_tool()]
-
-        # Initialize agent
-        agent = initialize_agent(
-            tools=tools,
-            llm=llm,
-            agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
-            memory=memory,
-            verbose=True,
-            handle_parsing_errors=True
-        )
-
-        return agent
-
-    except Exception as e:
-        logger.error(f"Error creating LangChain agent: {e}")
-        raise
-
 # --- UI: MathJax Configuration ---
 mathjax_config = '''
 <script>
@@ -473,24 +444,21 @@ def smart_truncate(text, max_length=3000):
     words = text[:max_length].split()
     return ' '.join(words[:-1]) + "... [Response truncated]"
 
-def generate_response_with_langchain(message, max_retries=3):
-    """Generate response using LangChain agent with proper message handling."""
+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()
             
-            # Initialize system prompt if not already done
-            initialize_system_prompt(current_agent)
-            
-            # Use the agent directly with the message
-            response = current_agent.run(input=message)
+            # 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"LangChain error (attempt {attempt + 1}): {e}")
+            logger.error(f"Agent error (attempt {attempt + 1}): {e}")
             if attempt < max_retries - 1:
                 time.sleep(2)
                 continue
@@ -520,8 +488,8 @@ def chat_response(message, history=None):
         except Exception as metrics_error:
             logger.error(f"Error in metrics_tracker.log_interaction: {metrics_error}")
         
-        # Generate response with LangChain
-        response = generate_response_with_langchain(message)
+        # Generate response with modern agent
+        response = generate_response_with_agent(message)
         
         # Log final metrics
         try:
@@ -541,7 +509,7 @@ def chat_response(message, history=None):
         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 - no voice parameter."""
+    """Main function to handle user submission."""
     if not message.strip():
         return history, ""
     
@@ -556,11 +524,10 @@ def respond_and_update(message, history):
     yield history, ""
 
 def clear_chat():
-    """Clear the chat history and reset system prompt flag."""
-    global agent, system_prompt_initialized
+    """Clear the chat history."""
+    global agent
     if agent is not None:
         agent.memory.clear()
-    system_prompt_initialized = False
     return [], ""
 
 def warmup_agent():
@@ -569,14 +536,11 @@ def warmup_agent():
     try:
         current_agent = get_agent()
         
-        # Initialize system prompt
-        initialize_system_prompt(current_agent)
-        
         # Run a simple test query
-        test_response = current_agent.run(input="Hello, this is a warmup test.")
+        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 the test interaction from memory
+        # Clear any test data from memory
         current_agent.memory.clear()
         
     except Exception as e:
@@ -640,7 +604,7 @@ def create_interface():
                     send = gr.Button("Send", elem_classes=["send-button"], size="sm")
                     clear = gr.Button("Clear", elem_classes=["clear-button"], size="sm")
             
-            # Event handlers - no voice parameter
+            # 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])
@@ -651,7 +615,6 @@ def create_interface():
         return demo
 
 # --- Main Execution ---
-# At the end of your app.py file, replace the main execution block:
 if __name__ == "__main__":
     try:
         logger.info("=" * 50)
@@ -661,7 +624,7 @@ if __name__ == "__main__":
         # Step 1: Preload the model and agent
         logger.info("Loading AI model...")
         start_time = time.time()
-        agent = create_langchain_agent()
+        agent = Educational_Agent()
         load_time = time.time() - start_time
         logger.info(f"Model loaded successfully in {load_time:.2f} seconds")