Update app.py

app.py

@@ -663,19 +663,92 @@ class Educational_Agent:
         tools = [Create_Graph_Tool]
         tool_node = ToolNode(tools)
         
-        def should_continue(state: EducationalAgentState) -> str:
-            """Determine next step in the workflow"""
+        def call_model(state: EducationalAgentState) -> dict:
+            """Call the LLM to generate a response"""
+            start_call_model_time = time.perf_counter()
+            current_time = datetime.now()
+            
+            messages = state["messages"]
+            
+            # Get the latest human message
+            user_query = ""
+            for msg in reversed(messages):
+                if isinstance(msg, HumanMessage):
+                    user_query = msg.content
+                    break
+            
+            if not user_query:
+                return {"messages": [AIMessage(content="I didn't receive a question. Please ask me something!")]}
+            
+            try:
+                # Check if tools are needed based on state
+                needs_tools = state.get("needs_tools", False)
+                
+                if needs_tools:
+                    # Create tool prompt for visualization
+                    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. 
+    
+    Format your response with explanatory text followed by a JSON block like this:
+    
+    ```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"
+    }}
+    ```
+    
+    Provide your educational response followed by the JSON configuration.
+    """
+                    prompt = tool_prompt
+                else:
+                    prompt = user_query
+                
+                # Generate response using the LLM (self.llm is accessible here)
+                response = self.llm.invoke(prompt)
+                
+                # Create AI message
+                ai_message = AIMessage(content=response)
+                
+                end_call_model_time = time.perf_counter()
+                call_model_time = end_call_model_time - start_call_model_time
+                log_metric(f"Call model time: {call_model_time:0.4f} seconds. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
+                
+                return {"messages": [ai_message]}
+                
+            except Exception as e:
+                logger.error(f"Error in call_model: {e}")
+                end_call_model_time = time.perf_counter()
+                call_model_time = end_call_model_time - start_call_model_time
+                log_metric(f"Call model time (error): {call_model_time:0.4f} seconds. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
+                
+                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"]
-            last_message = messages[-1]
             
-            # Check if we have tool calls to execute
-            if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
-                logger.info("Executing tools based on model decision")
-                return "tools"
+            # 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": []}
             
-            # Check if the message content contains JSON for tool calling
-            if isinstance(last_message, AIMessage) and last_message.content:
-                            content = last_message.content
+            content = last_ai_message.content
             
             # Look for JSON blocks in the message
             json_pattern = r'```json\s*\n?(.*?)\n?```'
@@ -721,7 +794,7 @@ class Educational_Agent:
             except Exception as e:
                 logger.error(f"Error processing JSON tools: {e}")
                 return {"messages": []}
-
+    
         def make_tool_decision(state: EducationalAgentState) -> dict:
             """Decide whether tools are needed and update state"""
             start_tool_decision_time = time.perf_counter()
@@ -747,7 +820,7 @@ class Educational_Agent:
             log_metric(f"Tool decision workflow time: {tool_decision_time:0.4f} seconds. Decision: {needs_visualization}. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
             
             return {"needs_tools": needs_visualization}
-
+    
         # Create the workflow graph
         workflow = StateGraph(EducationalAgentState)