Update app.py

app.py

@@ -131,76 +131,6 @@ hf_token = HF_TOKEN
 if not hf_token:
     logger.warning("Neither HF_TOKEN nor HUGGINGFACEHUB_API_TOKEN is set, the application may not work.")
 
-# LangGraph State Definition
-class EducationalAgentState(TypedDict):
-    messages: Annotated[Sequence[BaseMessage], add_messages]
-    needs_tools: bool
-    educational_context: Optional[str]
-
-@tool(return_direct=False)
-def Create_Graph_Tool(graph_config: str) -> str:
-    """
-    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.
-    """
-    start_create_graph_tool_time = time.perf_counter()
-    current_time = datetime.now()
-    
-    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>'
-            result = context_html + graph_html
-        else:
-            result = graph_html
-        
-        end_create_graph_tool_time = time.perf_counter()
-        graph_create_graph_tool_time = end_create_graph_tool_time - start_create_graph_tool_time
-        log_metric(f"Graph tool creation time: {graph_create_graph_tool_time:0.4f} seconds. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
-        
-        return result
-        
-    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 (Updated for LangGraph)
 class Tool_Decision_Engine:
     """Uses LLM to intelligently decide when visualization tools would be beneficial"""
@@ -279,6 +209,83 @@ Decision:"""
             log_metric(f"Tool decision time (error): {graph_decision_time:0.4f} seconds. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
             return False
 
+# LangGraph State Definition
+class EducationalAgentState(TypedDict):
+    messages: Annotated[Sequence[BaseMessage], add_messages]
+    needs_tools: bool
+    educational_context: Optional[str]
+
+@tool(return_direct=False)
+def Create_Graph_Tool(
+    data: dict,
+    plot_type: str,
+    title: str = "Generated Plot",
+    x_label: str = "",
+    y_label: str = "",
+    educational_context: str = ""
+) -> str:
+    """
+    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
+    
+    Args:
+        data: Dictionary with string keys and numeric values {"Category A": 25, "Category B": 40}
+        plot_type: "bar", "line", or "pie"
+        title: Title for the chart
+        x_label: X-axis label
+        y_label: Y-axis label
+        educational_context: Explanation of why this visualization helps learning
+    """
+    start_create_graph_tool_time = time.perf_counter()
+    current_time = datetime.now()
+    
+    try:
+        # Call the generate_plot function directly
+        content, artifact = generate_plot(
+            data=data,
+            plot_type=plot_type,
+            title=title,
+            x_label=x_label,
+            y_label=y_label
+        )
+        
+        # Check if there was an error
+        if "error" in artifact:
+            return f'<p style="color:red;">Graph generation failed: {artifact["error"]}</p>'
+        
+        # Convert the base64 image to HTML
+        base64_image = artifact["base64_image"]
+        
+        # Add educational context if provided
+        context_html = ""
+        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>'
+        
+        # Create the complete HTML with image
+        result = f"""{context_html}
+<div style="text-align: center; margin: 20px 0;">
+    <img src="data:image/png;base64,{base64_image}" 
+         style="max-width: 100%; height: auto; border-radius: 8px; box-shadow: 0 2px 10px rgba(0,0,0,0.1);" 
+         alt="{title}" />
+</div>"""
+        
+        end_create_graph_tool_time = time.perf_counter()
+        graph_create_graph_tool_time = end_create_graph_tool_time - start_create_graph_tool_time
+        log_metric(f"Graph tool creation time: {graph_create_graph_tool_time:0.4f} seconds. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
+        
+        return result
+        
+    except Exception as e:
+        logger.error(f"Error in graph generation: {e}")
+        return f'<p style="color:red;">Error creating graph: {str(e)}</p>'
+
 # System Prompt with ReAct Framework for Phi-3-mini
 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. You do so concisely, without excessive filler language or flowery content.
 
@@ -702,7 +709,7 @@ class Educational_Agent:
         
     def _create_langgraph_workflow(self):
         """Create the complete LangGraph workflow with improved tool calling"""
-        # Define tools
+        # Use the updated Create_Graph_Tool
         tools = [Create_Graph_Tool]
         tool_node = ToolNode(tools)
         
@@ -728,33 +735,40 @@ class Educational_Agent:
                 needs_tools = state.get("needs_tools", False)
                 
                 if needs_tools:
-                    # Create tool prompt for visualization
+                    # Create tool prompt that guides the model to use structured parameters
                     tool_prompt = f"""
     You are an educational AI assistant. The user has asked: "{user_query}"
     
-    This query would benefit from a visualization. Please provide a helpful educational response AND include a JSON configuration for creating a graph or chart. 
+    This query would benefit from a visualization. Please call the Create_Graph_Tool with appropriate structured parameters.
     
-    Format your response with explanatory text followed by a JSON block like this:
+    For the data parameter, create a meaningful dictionary with string keys and numeric values that illustrate the concept being discussed.
     
-    ```json
-    {{
-    "data": {{"Category 1": value1, "Category 2": value2}},
-    "plot_type": "bar|line|pie",
-    "title": "Descriptive Title",
-    "x_label": "X Axis Label", 
-    "y_label": "Y Axis Label",
-    "educational_context": "Explanation of why this visualization helps learning"
-    }}
-    ```
+    Choose the appropriate plot_type:
+    - "bar" for comparing categories or discrete data
+    - "line" for showing trends over time or continuous relationships  
+    - "pie" for showing parts of a whole or proportions
     
-    Provide your educational response followed by the JSON configuration.
+    Create a descriptive title and appropriate axis labels. Include an educational_context explaining why this visualization helps learning.
+    
+    Call the tool with these structured parameters, don't format as JSON.
     """
                     prompt = tool_prompt
                 else:
                     prompt = user_query
                 
-                # Generate response using the LLM (self.llm is accessible here)
-                response = self.llm.invoke(prompt)
+                # Bind tools to LLM if needed
+                if needs_tools:
+                    model_with_tools = self.llm
+                    # For Phi-3, we need to manually bind tools if supported
+                    try:
+                        if hasattr(self.llm, 'bind_tools'):
+                            model_with_tools = self.llm.bind_tools(tools)
+                        response = model_with_tools.invoke(prompt)
+                    except:
+                        # Fallback if tool binding not supported
+                        response = self.llm.invoke(prompt)
+                else:
+                    response = self.llm.invoke(prompt)
                 
                 # Create AI message
                 ai_message = AIMessage(content=response)
@@ -774,69 +788,15 @@ class Educational_Agent:
                 error_message = AIMessage(content=f"I encountered an error generating a response: {str(e)}")
                 return {"messages": [error_message]}
     
-        def process_json_tools(state: EducationalAgentState) -> dict:
-            """Process JSON tool configurations from the model response"""
-            start_process_tools_time = time.perf_counter()
-            current_time = datetime.now()
-            
-            messages = state["messages"]
-            
-            # Get the last AI message
-            last_ai_message = None
-            for msg in reversed(messages):
-                if isinstance(msg, AIMessage):
-                    last_ai_message = msg
-                    break
-            
-            if not last_ai_message or not last_ai_message.content:
-                return {"messages": []}
+        def should_continue(state: EducationalAgentState) -> str:
+            """Route to tools or end based on the last message"""
+            last_message = state["messages"][-1]
             
-            content = last_ai_message.content
-            
-            # Look for JSON blocks in the message
-            json_pattern = r'```json\s*\n?(.*?)\n?```'
-            json_matches = re.findall(json_pattern, content, re.DOTALL)
-            
-            if not json_matches:
-                logger.info("No JSON configuration found in message")
-                return {"messages": []}
-            
-            # Process the first JSON match
-            json_config = json_matches[0].strip()
-            
-            try:
-                # Validate JSON
-                config_dict = json.loads(json_config)
-                
-                # Check if it's a valid graph configuration
-                required_keys = ['data', 'plot_type', 'title']
-                if all(key in config_dict for key in required_keys):
-                    logger.info("Processing valid graph configuration")
-                    
-                    # Call the graph tool
-                    tool_result = Create_Graph_Tool.invoke({"graph_config": json_config})
-                    
-                    # Create a tool message
-                    tool_message = ToolMessage(
-                        content=tool_result,
-                        tool_call_id="graph_tool_call_1"
-                    )
-                    
-                    end_process_tools_time = time.perf_counter()
-                    process_tools_time = end_process_tools_time - start_process_tools_time
-                    log_metric(f"Process JSON tools time: {process_tools_time:0.4f} seconds. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
-                    
-                    return {"messages": [tool_message]}
-                else:
-                    logger.warning("JSON found but missing required graph configuration keys")
-                    return {"messages": []}
-                    
-            except json.JSONDecodeError as e:
-                logger.error(f"Invalid JSON in message: {e}")
-                return {"messages": []}
-            except Exception as e:
-                logger.error(f"Error processing JSON tools: {e}")
-                return {"messages": []}
+            # Check if the last message has tool calls
+            if hasattr(last_message, "tool_calls") and last_message.tool_calls:
+                return "tools"
+            else:
+                return END
     
         def make_tool_decision(state: EducationalAgentState) -> dict:
             """Decide whether tools are needed and update state"""
@@ -870,13 +830,21 @@ class Educational_Agent:
         # Add nodes
         workflow.add_node("decide_tools", make_tool_decision)
         workflow.add_node("call_model", call_model)
-        workflow.add_node("process_tools", process_json_tools)
+        workflow.add_node("tools", tool_node)
         
         # Add edges
         workflow.add_edge(START, "decide_tools")
         workflow.add_edge("decide_tools", "call_model")
-        workflow.add_edge("call_model", "process_tools")
-        workflow.add_edge("process_tools", END)
+        
+        # Add conditional edge from call_model
+        workflow.add_conditional_edges(
+            "call_model",
+            should_continue,
+            {"tools": "tools", END: END}
+        )
+        
+        # After tools, go back to call_model for final response
+        workflow.add_edge("tools", "call_model")
         
         # Compile the workflow
         return workflow.compile(checkpointer=MemorySaver())