Update app.py

app.py

@@ -832,7 +832,93 @@ class ResponsePostProcessor:
 
 post_processor = ResponsePostProcessor()
 
+# ============================================================================
+# DATA EXTRACTION FOR GRAPHING 
+# ============================================================================
+
+def extract_graph_data(user_input: str, conversation_history: List[Dict]) -> Optional[Dict]:
+    """
+    Use LLM to extract graphable data from user input.
+    
+    Returns:
+        Dict with keys: data, plot_type, title, x_label, y_label, educational_context
+        Or None if no data can be extracted
+    """
+    from model_manager import get_model
+    
+    model = get_model()
+    
+    # Format conversation context
+    context = ""
+    if conversation_history:
+        recent = conversation_history[-2:]
+        context = "\n".join([f"{msg['role']}: {msg['content'][:200]}" for msg in recent])
+    
+    extraction_prompt = f"""Extract graphable data from the user's message.
+
+Previous context:
+{context}
+
+Current message: {user_input}
+
+If the message contains data that can be graphed (numbers, comparisons, datasets, trends), extract:
+1. The data as key-value pairs
+2. The best plot type (bar, line, or pie)
+3. A descriptive title
+4. Axis labels (if applicable)
+5. Educational context explaining what the graph shows
+
+Respond in JSON format ONLY:
+{{
+    "has_data": true/false,
+    "data": {{"label1": value1, "label2": value2, ...}},
+    "plot_type": "bar/line/pie",
+    "title": "Graph Title",
+    "x_label": "X Axis Label",
+    "y_label": "Y Axis Label",
+    "educational_context": "Brief explanation of what this graph represents"
+}}
+
+If no graphable data exists, respond: {{"has_data": false}}"""
+
+    try:
+        system_prompt = "You are a data extraction expert. Extract graphable data from text and respond in valid JSON only."
+        
+        response = model.generate(
+            system_prompt=system_prompt,
+            user_message=extraction_prompt,
+            max_tokens=300,
+            temperature=0.3
+        )
+        
+        # Parse JSON response
+        import json
+        # Try to extract JSON from response
+        json_start = response.find('{')
+        json_end = response.rfind('}') + 1
+        
+        if json_start == -1 or json_end == 0:
+            return None
+            
+        json_str = response[json_start:json_end]
+        result = json.loads(json_str)
+        
+        if not result.get('has_data', False):
+            return None
+            
+        # Validate required fields
+        if 'data' not in result or not result['data']:
+            return None
+            
+        return result
+        
+    except Exception as e:
+        logger.error(f"Data extraction failed: {e}")
+        return None
+
 
+
+        
 # ============================================================================
 # TOOL FUNCTIONS
 # ============================================================================
@@ -892,12 +978,11 @@ def orchestrate_turn(user_input: str, session_id: str = "default") -> str:
     """
     Main orchestration function implementing the redesign workflow.
     
-    OPTIMIZED: Uses single Qwen3-Claude GGUF (loads once, all agents share)
     
     Steps:
     1. Reset prompt state
     2. Process user input (history)
-    3. Tool decision (Qwen3-Claude)
+    3. Tool decision
     4. Regex checks
     5. Agent execution (Qwen3-Claude)
     6. Thinking agents (Qwen3-Claude)
@@ -936,9 +1021,11 @@ def orchestrate_turn(user_input: str, session_id: str = "default") -> str:
         log_step("Step 2: Process user input", step_start)
         
         # ====================================================================
-        # STEP 3: TOOL DECISION ENGINE (Llama-3.2-3B)
+        # STEP 3: TOOL DECISION ENGINE & GRAPH GENERATION
         # ====================================================================
-        step_start = log_step("Step 3: Tool decision")
+        step_start = log_step("Step 3: Tool decision & graph generation")
+        
+        # Check if visualization is needed
         tool_decision_result = tool_agent.decide(user_input, conversation_history)
         
         tool_img_output = ""
@@ -947,10 +1034,35 @@ def orchestrate_turn(user_input: str, session_id: str = "default") -> str:
         if tool_decision_result:
             logger.info("Tool decision: YES - visualization needed")
             prompt_state.update("TOOL_USE_ENHANCEMENT", True)
+            
+            # Extract data and generate graph
+            logger.info("→ Extracting graphable data from input...")
+            graph_data = extract_graph_data(user_input, conversation_history)
+            
+            if graph_data:
+                logger.info(f"✓ Data extracted: {len(graph_data['data'])} data points")
+                logger.info(f"  Plot type: {graph_data['plot_type']}")
+                
+                # Generate the graph
+                tool_img_output = Create_Graph_Tool(
+                    data=graph_data['data'],
+                    plot_type=graph_data['plot_type'],
+                    title=graph_data.get('title', 'Generated Plot'),
+                    x_label=graph_data.get('x_label', ''),
+                    y_label=graph_data.get('y_label', ''),
+                    educational_context=graph_data.get('educational_context', '')
+                )
+                
+                # Store context for thinking agents
+                tool_context = f"Graph created: {graph_data['plot_type']} chart showing {graph_data.get('title', 'data visualization')}"
+                
+                logger.info("✓ Graph generated successfully")
+            else:
+                logger.info("⚠️ No extractable data found - skipping graph generation")
         else:
             logger.info("Tool decision: NO - no visualization needed")
         
-        log_step("Step 3: Tool decision", step_start)
+        log_step("Step 3: Tool decision & graph generation", step_start)
         # ====================================================================
         # STEP 4: REGEX LOGICAL EXPRESSIONS
         # ====================================================================