Applied changes to resolve noted issues.

https://github.com/Jdesiree112/Prompt_Engineering/blob/main/CaseStudy_Mimir/Issue_ProblemSolvingCaseStudy/Issue-Solutions_Overview.md

app.py

@@ -3,8 +3,8 @@ from graph_tool import generate_plot
 from metrics import MimirMetrics
 import os
 
+# Updated environment variables
 os.environ['HF_HOME'] = '/tmp/huggingface'
-os.environ['TRANSFORMERS_CACHE'] = '/tmp/huggingface'
 os.environ['HF_DATASETS_CACHE'] = '/tmp/huggingface'
 
 import time
@@ -14,7 +14,8 @@ 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 langchain_core.runnables import RunnableBranch, Runnable
+from langchain_core.runnables.utils import Input, Output
 from langchain.agents import initialize_agent, AgentType
 from langchain.memory import ConversationBufferWindowMemory
 from typing import Optional, List, Any, Type
@@ -103,28 +104,42 @@ class Tool_Decision_Engine:
         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
+
+EXCLUDE if query is:
+- Greetings or casual conversation (hello, hi, hey)
+- Simple definitions without data
+- Test/warmup messages
+- General explanations that don't involve data
+
+INCLUDE if query involves:
+- Mathematical functions or relationships
+- Data analysis or statistics
+- Comparisons that benefit from charts
+- Trends or patterns over time
+- Creating practice questions with 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"""
+        """Enhanced decision logic with explicit exclusions"""
         try:
+            # Explicit exclusions for common non-visual queries
+            exclusion_patterns = [
+                r'^(hello|hi|hey)\b',
+                r'warmup.*test',
+                r'(what is|define|explain)\s+\w+\s*(of|the)?',
+                r'capital\s+of',
+                r'^(greet|greeting)'
+            ]
+            
+            query_lower = query.lower().strip()
+            
+            # Check exclusions first
+            for pattern in exclusion_patterns:
+                if re.search(pattern, query_lower):
+                    return False
+            
             # Create decision prompt
             decision_query = self.decision_prompt.format(query=query)
             
@@ -137,7 +152,10 @@ Decision:"""
             # 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
+            # More strict parsing
+            if "YES" in decision_text and "NO" not in decision_text:
+                return True
+            return False
             
         except Exception as e:
             logger.error(f"Error in tool decision making: {e}")
@@ -191,11 +209,14 @@ You have the ability to create graphs and charts to enhance your explanations. U
 - 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 ---
+# --- Fixed LLM Class with Runnable inheritance ---
 logger = logging.getLogger(__name__)
 
-class Qwen25SmallLLM:
+class Qwen25SmallLLM(Runnable):
+    """LLM class that properly inherits from Runnable for LangChain compatibility"""
+    
     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:
@@ -212,12 +233,12 @@ class Qwen25SmallLLM:
                     llm_int8_skip_modules=["lm_head"]
                 )
 
-                # Try quantized load
+                # Try quantized load with updated dtype parameter
                 self.model = AutoModelForCausalLM.from_pretrained(
                     model_path,
                     quantization_config=quant_config,
                     device_map="auto",
-                    torch_dtype=torch.bfloat16,
+                    dtype=torch.bfloat16,
                     trust_remote_code=True,
                     low_cpu_mem_usage=True
                 )
@@ -236,13 +257,20 @@ class Qwen25SmallLLM:
         """Fallback if quantization fails."""
         self.model = AutoModelForCausalLM.from_pretrained(
             model_path,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            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:
+    def invoke(self, input: Input, config=None) -> Output:
+        """Main invoke method for Runnable compatibility"""
+        # Handle both string and dict inputs for flexibility
+        if isinstance(input, dict):
+            prompt = input.get('input', str(input))
+        else:
+            prompt = str(input)
+            
         try:
             messages = [
                 {"role": "system", "content": SYSTEM_PROMPT},
@@ -273,6 +301,14 @@ class Qwen25SmallLLM:
             logger.error(f"Generation error: {e}")
             return f"[Error generating response: {str(e)}]"
 
+    @property
+    def InputType(self) -> Type[Input]:
+        return str
+
+    @property 
+    def OutputType(self) -> Type[Output]:
+        return str
+
 # --- Modern Agent Implementation ---
 class Educational_Agent:
     """Modern LangChain agent with LLM-based tool decision making"""
@@ -633,7 +669,7 @@ if __name__ == "__main__":
         
         # Step 2: Warm up the model
         logger.info("Warming up model...")
-        warmup_agent()  # Comment out temporarily to avoid hanging
+        warmup_agent()
         
         interface = create_interface()
         interface.queue()