Update app.py

app.py

@@ -15,7 +15,7 @@ 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.agents import initialize_agent, AgentType
 from langchain.memory import ConversationBufferWindowMemory
 from typing import Optional, List, Any, Type
 from pydantic import BaseModel, Field
@@ -102,30 +102,24 @@ class Tool_Decision_Engine:
         self.conversation_history = []
         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:
@@ -152,13 +146,11 @@ Decision:"""
 
 # --- 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
@@ -168,16 +160,13 @@ SYSTEM_PROMPT = """You are Mimir, an expert multi-concept tutor designed to faci
 - 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)
@@ -187,15 +176,12 @@ You have the ability to create graphs and charts to enhance your explanations. U
 - 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
@@ -203,7 +189,6 @@ You have the ability to create graphs and charts to enhance your explanations. U
 - 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 ---
@@ -225,17 +210,19 @@ class Qwen25SmallLLM:
                     bnb_4bit_quant_type="nf4", 
                     llm_int8_threshold=0.0,
                     llm_int8_skip_modules=["lm_head"]
-                    )
-
-            # 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
-            )
+                )
+
+                # 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
+                )
+            else:
+                self._load_fallback_model(model_path)
 
         except Exception as e:
             logger.warning(f"Quantized load failed, falling back: {e}")
@@ -294,7 +281,6 @@ class Educational_Agent:
         self.llm = Qwen25SmallLLM(model_path="Qwen/Qwen2.5-1.5B-Instruct", use_4bit=True)
         self.tool_decision_engine = Tool_Decision_Engine(self.llm)
         self.conversation_history = []
-        )
         
     def should_use_tools(self, query: str) -> bool:
         """Use LLM reasoning to determine if tools are needed"""
@@ -313,66 +299,81 @@ class Educational_Agent:
         ])
     
     def process_with_tools(self, query: str) -> str:
-   """Process query with tools available using traditional agent"""
-   try:
-       # Create agent with tools using traditional approach
-       tools = [Create_Graph_Tool]
-       
-       # Create memory for this agent instance
-       memory = ConversationBufferWindowMemory(
-           memory_key="chat_history",
-           return_messages=True
-       )
-       
-       # Add conversation history to memory
-       for exchange in self.conversation_history:
-           if exchange.startswith("User: "):
-               memory.chat_memory.add_user_message(exchange[6:])
-           elif exchange.startswith("Assistant: "):
-               memory.chat_memory.add_ai_message(exchange[11:])
-       
-       # Use traditional agent with better prompting
-       agent = initialize_agent(
-           tools=tools,
-           llm=self.llm,
-           agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-           memory=memory,
-           verbose=True,
-           handle_parsing_errors=True,
-           agent_kwargs={
-               "prefix": SYSTEM_PROMPT + "\n\nYou have access to graph creation tools. Use them when visualization would genuinely help explain concepts. Think carefully about whether a tool is needed before using it.\n\n"
-           }
-       )
-       
-       response = agent.run(query)
-       
-       # Store the exchange
-       self.conversation_history.append(f"User: {query}")
-       self.conversation_history.append(f"Assistant: {response}")
-       
-       return 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)}"
+        """Process query with tools available using traditional agent"""
+        try:
+            # Create agent with tools using traditional approach
+            tools = [Create_Graph_Tool]
+            
+            # Create memory for this agent instance
+            memory = ConversationBufferWindowMemory(
+                memory_key="chat_history",
+                return_messages=True
+            )
+            
+            # Add conversation history to memory
+            for exchange in self.conversation_history:
+                if exchange.startswith("User: "):
+                    memory.chat_memory.add_user_message(exchange[6:])
+                elif exchange.startswith("Assistant: "):
+                    memory.chat_memory.add_ai_message(exchange[11:])
+            
+            # Use traditional agent with better prompting
+            agent = initialize_agent(
+                tools=tools,
+                llm=self.llm,
+                agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+                memory=memory,
+                verbose=True,
+                handle_parsing_errors=True,
+                agent_kwargs={
+                    "prefix": SYSTEM_PROMPT + "\n\nYou have access to graph creation tools. Use them when visualization would genuinely help explain concepts. Think carefully about whether a tool is needed before using it.\n\n"
+                }
+            )
+            
+            response = agent.run(query)
+            
+            # Store the exchange
+            self.conversation_history.append(f"User: {query}")
+            self.conversation_history.append(f"Assistant: {response}")
+            
+            return 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:
-       # Include recent conversation in prompt
-       context = "\n".join(self.conversation_history[-8:]) if self.conversation_history else ""
-       full_prompt = f"Recent conversation:\n{context}\n\nCurrent query: {query}" if context else query
-       
-       response = self.llm.invoke(full_prompt)
-       
-       # Store the exchange
-       self.conversation_history.append(f"User: {query}")
-       self.conversation_history.append(f"Assistant: {response}")
-       
-       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)}"
+        """Process query without tools"""
+        try:
+            # Include recent conversation in prompt
+            context = "\n".join(self.conversation_history[-8:]) if self.conversation_history else ""
+            full_prompt = f"Recent conversation:\n{context}\n\nCurrent query: {query}" if context else query
+            
+            response = self.llm.invoke(full_prompt)
+            
+            # Store the exchange
+            self.conversation_history.append(f"User: {query}")
+            self.conversation_history.append(f"Assistant: {response}")
+            
+            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
@@ -632,7 +633,7 @@ if __name__ == "__main__":
         
         # Step 2: Warm up the model
         logger.info("Warming up model...")
-        warmup_agent()
+        warmup_agent()  # Comment out temporarily to avoid hanging
         
         interface = create_interface()
         interface.queue()