Update app.py

app.py

@@ -182,16 +182,17 @@ def discover_attachment(task_id: str, api_url: str) -> Optional[str]:
 # -------------------------
 
 class AgentState(TypedDict):
-    question: Annotated[str, override]
-    current_step: Annotated[str, override]
-    final_answer: Annotated[str, override]
-    history: Annotated[List[Dict[str, str]], list.__add__]
-    needs_search: Annotated[bool, override]
-    search_query: Annotated[str, override]
-    task_id: Annotated[str, override]
-    logs: Annotated[Dict[str, Any], merge_dicts]
-    file_url: Annotated[str, override]
-    code_blocks: Annotated[List[Dict[str, str]], list.__add__]
+    question: str
+    current_step: str
+    final_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  # Add this field for routing
 
 # -------------------------
 # BasicAgent implementation
@@ -247,7 +248,7 @@ class BasicAgent:
 
         state: AgentState = {
             "question": question,
-            "current_step": "answer",
+            "current_step": "route",
             "final_answer": "",
             "history": [],
             "needs_search": False,
@@ -255,7 +256,8 @@ class BasicAgent:
             "task_id": task_id,
             "logs": {},
             "file_url": file_url,
-            "code_blocks": []
+            "code_blocks": [],
+            "next_step": ""  # Initialize next_step
         }
         
         print(f"\nProcessing task {task_id}")
@@ -265,7 +267,47 @@ class BasicAgent:
         final_state = self.workflow.invoke(state)
         return final_state["final_answer"]
 
-    def _process_image(self, state: AgentState) -> AgentState:
+    def _route_to_tool(self, state: AgentState) -> Dict[str, Any]:
+        """Route the state to the appropriate tool based on file type."""
+        if not state["file_url"]:
+            return {"next_step": "process_text"}
+            
+        try:
+            response = SESSION.get(state["file_url"], timeout=30)
+            response.raise_for_status()
+            data = response.content
+            
+            # Get content type from response headers first, fallback to URL-based detection
+            kind = response.headers.get("Content-Type", "")
+            if kind in ("application/octet-stream", ""):
+                # rough sniff: look at the first few bytes
+                sig = data[:4]
+                if sig.startswith(b"\x89PNG"):
+                    kind = "image/png"
+                elif sig.startswith(b"\xFF\xD8"):
+                    kind = "image/jpeg"
+                elif sig[:2] == b"PK":  # XLSX = ZIP
+                    kind = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+            elif not kind:  # fallback if header missing
+                kind = mimetypes.guess_type(state["file_url"])[0] or ""
+            
+            if "image" in kind:
+                return {"next_step": "process_image"}
+            elif "video" in kind:
+                return {"next_step": "process_video"}
+            elif "spreadsheet" in kind or "excel" in kind:
+                return {"next_step": "process_spreadsheet"}
+            elif state["file_url"].endswith(".py"):
+                return {"next_step": "process_python"}
+            else:
+                print(f"Unsupported file type: {kind}")
+                return {"next_step": "process_text"}
+                
+        except Exception as e:
+            print(f"Error determining file type: {str(e)}")
+            return {"next_step": "process_text"}
+
+    def _process_image(self, state: AgentState) -> Dict[str, Any]:
         """Process image files using LLaVA."""
         try:
             print(f"Downloading {state['file_url']} …")
@@ -278,16 +320,20 @@ class BasicAgent:
             answer = image_qa_bytes(data, state["question"])
             
             print(f"Generated answer: {answer}")
-            state["final_answer"] = answer
-            state["current_step"] = "done"
-            return state
+            return {
+                "final_answer": answer,
+                "current_step": "done",
+                "next_step": "done"
+            }
         except Exception as e:
             print(f"\nError processing image {state['file_url']}: {str(e)}")
-            state["final_answer"] = f"Error processing image: {str(e)}"
-            state["current_step"] = "done"
-            return state
+            return {
+                "final_answer": f"Error processing image: {str(e)}",
+                "current_step": "done",
+                "next_step": "done"
+            }
 
-    def _process_video(self, state: AgentState) -> AgentState:
+    def _process_video(self, state: AgentState) -> Dict[str, Any]:
         """Process video files using VideoMAE."""
         try:
             print(f"Downloading {state['file_url']} …")
@@ -300,16 +346,20 @@ class BasicAgent:
             answer = video_label_bytes(data)
             
             print(f"Generated answer: {answer}")
-            state["final_answer"] = answer
-            state["current_step"] = "done"
-            return state
+            return {
+                "final_answer": answer,
+                "current_step": "done",
+                "next_step": "done"
+            }
         except Exception as e:
             print(f"\nError processing video {state['file_url']}: {str(e)}")
-            state["final_answer"] = f"Error processing video: {str(e)}"
-            state["current_step"] = "done"
-            return state
+            return {
+                "final_answer": f"Error processing video: {str(e)}",
+                "current_step": "done",
+                "next_step": "done"
+            }
 
-    def _process_spreadsheet(self, state: AgentState) -> AgentState:
+    def _process_spreadsheet(self, state: AgentState) -> Dict[str, Any]:
         """Process spreadsheet files."""
         try:
             print(f"Downloading {state['file_url']} …")
@@ -322,16 +372,20 @@ class BasicAgent:
             answer = sheet_answer_bytes(data)
             
             print(f"Generated answer: {answer}")
-            state["final_answer"] = answer
-            state["current_step"] = "done"
-            return state
+            return {
+                "final_answer": answer,
+                "current_step": "done",
+                "next_step": "done"
+            }
         except Exception as e:
             print(f"\nError processing spreadsheet {state['file_url']}: {str(e)}")
-            state["final_answer"] = f"Error processing spreadsheet: {str(e)}"
-            state["current_step"] = "done"
-            return state
+            return {
+                "final_answer": f"Error processing spreadsheet: {str(e)}",
+                "current_step": "done",
+                "next_step": "done"
+            }
 
-    def _process_python(self, state: AgentState) -> AgentState:
+    def _process_python(self, state: AgentState) -> Dict[str, Any]:
         """Process Python files."""
         try:
             print(f"Downloading {state['file_url']} …")
@@ -344,16 +398,20 @@ class BasicAgent:
             answer = run_python(data.decode())
             
             print(f"Generated answer: {answer}")
-            state["final_answer"] = answer
-            state["current_step"] = "done"
-            return state
+            return {
+                "final_answer": answer,
+                "current_step": "done",
+                "next_step": "done"
+            }
         except Exception as e:
             print(f"\nError processing Python file {state['file_url']}: {str(e)}")
-            state["final_answer"] = f"Error processing Python file: {str(e)}"
-            state["current_step"] = "done"
-            return state
+            return {
+                "final_answer": f"Error processing Python file: {str(e)}",
+                "current_step": "done",
+                "next_step": "done"
+            }
 
-    def _process_text(self, state: AgentState) -> AgentState:
+    def _process_text(self, state: AgentState) -> Dict[str, Any]:
         """Process text-only questions using LLM."""
         print("\nProcessing as text-only question...")
         prompt = f"""
@@ -368,53 +426,18 @@ QUESTION:
             raw = self._call_llm(prompt, 300)
             answer = self._safe_parse(raw)
             print(f"Generated answer: {answer}")
-            state["final_answer"] = answer
+            return {
+                "final_answer": answer,
+                "current_step": "done",
+                "next_step": "done"
+            }
         except Exception as e:
             print(f"\nLLM Error in answer generation: {str(e)}")
-            state["final_answer"] = "I encountered an error while generating the answer."
-        
-        state["current_step"] = "done"
-        return state
-
-    def _route_to_tool(self, state: AgentState) -> str:
-        """Route the state to the appropriate tool based on file type."""
-        if not state["file_url"]:
-            return "process_text"
-            
-        try:
-            response = SESSION.get(state["file_url"], timeout=30)
-            response.raise_for_status()
-            data = response.content
-            
-            # Get content type from response headers first, fallback to URL-based detection
-            kind = response.headers.get("Content-Type", "")
-            if kind in ("application/octet-stream", ""):
-                # rough sniff: look at the first few bytes
-                sig = data[:4]
-                if sig.startswith(b"\x89PNG"):
-                    kind = "image/png"
-                elif sig.startswith(b"\xFF\xD8"):
-                    kind = "image/jpeg"
-                elif sig[:2] == b"PK":  # XLSX = ZIP
-                    kind = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
-            elif not kind:  # fallback if header missing
-                kind = mimetypes.guess_type(state["file_url"])[0] or ""
-            
-            if "image" in kind:
-                return "process_image"
-            elif "video" in kind:
-                return "process_video"
-            elif "spreadsheet" in kind or "excel" in kind:
-                return "process_spreadsheet"
-            elif state["file_url"].endswith(".py"):
-                return "process_python"
-            else:
-                print(f"Unsupported file type: {kind}")
-                return "process_text"
-                
-        except Exception as e:
-            print(f"Error determining file type: {str(e)}")
-            return "process_text"
+            return {
+                "final_answer": "I encountered an error while generating the answer.",
+                "current_step": "done",
+                "next_step": "done"
+            }
 
     def _build_workflow(self) -> Graph:
         """Build the workflow graph with conditional edges."""
@@ -431,15 +454,15 @@ QUESTION:
         # Set entry point
         sg.set_entry_point("route")
         
-        # Add conditional edges
+        # Add conditional edges based on next_step field
         sg.add_conditional_edges(
             "route",
             {
-                "process_image": lambda x: x == "process_image",
-                "process_video": lambda x: x == "process_video",
-                "process_spreadsheet": lambda x: x == "process_spreadsheet",
-                "process_python": lambda x: x == "process_python",
-                "process_text": lambda x: x == "process_text"
+                "process_image": lambda x: x["next_step"] == "process_image",
+                "process_video": lambda x: x["next_step"] == "process_video",
+                "process_spreadsheet": lambda x: x["next_step"] == "process_spreadsheet",
+                "process_python": lambda x: x["next_step"] == "process_python",
+                "process_text": lambda x: x["next_step"] == "process_text"
             }
         )