Update app.py

app.py

@@ -182,17 +182,12 @@ def discover_attachment(task_id: str, api_url: str) -> Optional[str]:
 # -------------------------
 
 class AgentState(TypedDict):
-    question: str
-    current_step: str
-    answer: str
-    history: List[Dict[str, str]]
-    needs_search: bool
-    search_query: str
-    task_id: str
-    logs: Dict[str, Any]
-    file_url: str
-    code_blocks: List[Dict[str, str]]
-    next_step: str
+    question: str            # Input question
+    answer: str             # Output answer (required by Gradio)
+    current_step: str       # Current processing step
+    next_step: str         # Next step in workflow
+    file_url: str          # URL of attached file if any
+    history: List[Dict[str, str]]  # Conversation history
 
 # -------------------------
 # BasicAgent implementation
@@ -246,18 +241,14 @@ class BasicAgent:
         if not file_url:
             file_url = discover_attachment(task_id, self.api_url) or ""
 
+        # Initialize state with just the question
         state: AgentState = {
             "question": question,
-            "current_step": "route",
             "answer": "",
-            "history": [],
-            "needs_search": False,
-            "search_query": "",
-            "task_id": task_id,
-            "logs": {},
+            "current_step": "route",
+            "next_step": "",
             "file_url": file_url,
-            "code_blocks": [],
-            "next_step": ""
+            "history": []
         }
         
         print(f"\nProcessing task {task_id}")
@@ -265,10 +256,15 @@ class BasicAgent:
         print(f"File URL: {state['file_url']}")
         
         try:
-            final_state = self.workflow.invoke(state)
+            # Invoke the workflow with just the question
+            final_state = self.workflow.invoke({"question": question})
+            
+            # Debug guard to check for answer key
             if "answer" not in final_state:
-                raise ValueError(f"No 'answer' in state: keys are {list(final_state.keys())}")
+                raise ValueError(f"☠   No 'answer' key in state; keys = {list(final_state.keys())}")
+                
             return final_state["answer"]
+            
         except Exception as e:
             print(f"Error in workflow execution: {str(e)}")
             return f"Error processing question: {str(e)}"
@@ -331,17 +327,13 @@ class BasicAgent:
             print(f"Generated answer: {answer}")
             return {
                 "answer": answer,
-                "current_step": "done",
-                "next_step": "done",
-                "history": state["history"] + [{"role": "assistant", "content": answer}]
+                "next_step": END
             }
         except Exception as e:
             print(f"\nError processing image {state['file_url']}: {str(e)}")
             return {
                 "answer": f"Error processing image: {str(e)}",
-                "current_step": "done",
-                "next_step": "done",
-                "history": state["history"] + [{"role": "assistant", "content": f"Error: {str(e)}"}]
+                "next_step": END
             }
 
     def _process_video(self, state: AgentState) -> Dict[str, Any]:
@@ -359,17 +351,13 @@ class BasicAgent:
             print(f"Generated answer: {answer}")
             return {
                 "answer": answer,
-                "current_step": "done",
-                "next_step": "done",
-                "history": state["history"] + [{"role": "assistant", "content": answer}]
+                "next_step": END
             }
         except Exception as e:
             print(f"\nError processing video {state['file_url']}: {str(e)}")
             return {
                 "answer": f"Error processing video: {str(e)}",
-                "current_step": "done",
-                "next_step": "done",
-                "history": state["history"] + [{"role": "assistant", "content": f"Error: {str(e)}"}]
+                "next_step": END
             }
 
     def _process_spreadsheet(self, state: AgentState) -> Dict[str, Any]:
@@ -387,17 +375,13 @@ class BasicAgent:
             print(f"Generated answer: {answer}")
             return {
                 "answer": answer,
-                "current_step": "done",
-                "next_step": "done",
-                "history": state["history"] + [{"role": "assistant", "content": answer}]
+                "next_step": END
             }
         except Exception as e:
             print(f"\nError processing spreadsheet {state['file_url']}: {str(e)}")
             return {
                 "answer": f"Error processing spreadsheet: {str(e)}",
-                "current_step": "done",
-                "next_step": "done",
-                "history": state["history"] + [{"role": "assistant", "content": f"Error: {str(e)}"}]
+                "next_step": END
             }
 
     def _process_python(self, state: AgentState) -> Dict[str, Any]:
@@ -415,17 +399,13 @@ class BasicAgent:
             print(f"Generated answer: {answer}")
             return {
                 "answer": answer,
-                "current_step": "done",
-                "next_step": "done",
-                "history": state["history"] + [{"role": "assistant", "content": answer}]
+                "next_step": END
             }
         except Exception as e:
             print(f"\nError processing Python file {state['file_url']}: {str(e)}")
             return {
                 "answer": f"Error processing Python file: {str(e)}",
-                "current_step": "done",
-                "next_step": "done",
-                "history": state["history"] + [{"role": "assistant", "content": f"Error: {str(e)}"}]
+                "next_step": END
             }
 
     def _process_text(self, state: AgentState) -> Dict[str, Any]:
@@ -445,17 +425,13 @@ QUESTION:
             print(f"Generated answer: {answer}")
             return {
                 "answer": answer,
-                "current_step": "done",
-                "next_step": "done",
-                "history": state["history"] + [{"role": "assistant", "content": answer}]
+                "next_step": END
             }
         except Exception as e:
             print(f"\nLLM Error in answer generation: {str(e)}")
             return {
                 "answer": "I encountered an error while generating the answer.",
-                "current_step": "done",
-                "next_step": "done",
-                "history": state["history"] + [{"role": "assistant", "content": "Error generating answer"}]
+                "next_step": END
             }
 
     def _build_workflow(self) -> Graph:
@@ -473,7 +449,7 @@ QUESTION:
         # Set entry point
         sg.set_entry_point("route")
         
-        # Add conditional edges based on next_step field
+        # Add conditional edges from route to processing nodes
         sg.add_conditional_edges(
             "route",
             {
@@ -487,7 +463,7 @@ QUESTION:
         
         # Add edges from each processing node to END
         for node in ["process_image", "process_video", "process_spreadsheet", "process_python", "process_text"]:
-            sg.add_edge(node, END)
+            sg.add_edge(node, END)  # Critical: ensure each processing node terminates at END
             
         return sg.compile()