Update app.py

app.py

@@ -265,62 +265,96 @@ class BasicAgent:
         final_state = self.workflow.invoke(state)
         return final_state["final_answer"]
 
-    def _generate_answer(self, state: AgentState) -> AgentState:
-        if state["file_url"]:
-            try:
-                print(f"Downloading {state['file_url']} …")
-                response = SESSION.get(state["file_url"], timeout=30)
-                response.raise_for_status()
-                data = response.content
-                print(f"Successfully downloaded file, size: {len(data)} bytes")
-                
-                # 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 ""
-                print(f"Detected file type: {kind}")
-                
-                if "image" in kind:
-                    print("Processing as image...")
-                    answer = image_qa_bytes(data, state["question"])
-                elif "video" in kind:
-                    print("Processing as video...")
-                    answer = video_label_bytes(data)
-                elif "spreadsheet" in kind or "excel" in kind:
-                    print("Processing as spreadsheet...")
-                    answer = sheet_answer_bytes(data)
-                elif state["file_url"].endswith(".py"):
-                    print("Processing as Python file...")
-                    answer = run_python(data.decode())
-                else:
-                    print(f"Unsupported file type: {kind}")
-                    answer = f"Unsupported file type: {kind}"
-                
-                print(f"Generated answer: {answer}")
-                state["final_answer"] = answer
-                state["current_step"] = "done"
-                return state
-            except requests.exceptions.RequestException as e:
-                print(f"Error downloading file: {e}")
-                state["final_answer"] = f"Error downloading file: {str(e)}"
-                state["current_step"] = "done"
-                return state
-            except Exception as e:
-                print(f"\nError processing file {state['file_url']}: {str(e)}")
-                state["final_answer"] = f"Error processing file: {str(e)}"
-                state["current_step"] = "done"
-                return state
+    def _process_image(self, state: AgentState) -> AgentState:
+        """Process image files using LLaVA."""
+        try:
+            print(f"Downloading {state['file_url']} …")
+            response = SESSION.get(state["file_url"], timeout=30)
+            response.raise_for_status()
+            data = response.content
+            print(f"Successfully downloaded file, size: {len(data)} bytes")
+            
+            print("Processing as image...")
+            answer = image_qa_bytes(data, state["question"])
+            
+            print(f"Generated answer: {answer}")
+            state["final_answer"] = answer
+            state["current_step"] = "done"
+            return state
+        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
+
+    def _process_video(self, state: AgentState) -> AgentState:
+        """Process video files using VideoMAE."""
+        try:
+            print(f"Downloading {state['file_url']} …")
+            response = SESSION.get(state["file_url"], timeout=30)
+            response.raise_for_status()
+            data = response.content
+            print(f"Successfully downloaded file, size: {len(data)} bytes")
+            
+            print("Processing as video...")
+            answer = video_label_bytes(data)
+            
+            print(f"Generated answer: {answer}")
+            state["final_answer"] = answer
+            state["current_step"] = "done"
+            return state
+        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
+
+    def _process_spreadsheet(self, state: AgentState) -> AgentState:
+        """Process spreadsheet files."""
+        try:
+            print(f"Downloading {state['file_url']} …")
+            response = SESSION.get(state["file_url"], timeout=30)
+            response.raise_for_status()
+            data = response.content
+            print(f"Successfully downloaded file, size: {len(data)} bytes")
+            
+            print("Processing as spreadsheet...")
+            answer = sheet_answer_bytes(data)
+            
+            print(f"Generated answer: {answer}")
+            state["final_answer"] = answer
+            state["current_step"] = "done"
+            return state
+        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
+
+    def _process_python(self, state: AgentState) -> AgentState:
+        """Process Python files."""
+        try:
+            print(f"Downloading {state['file_url']} …")
+            response = SESSION.get(state["file_url"], timeout=30)
+            response.raise_for_status()
+            data = response.content
+            print(f"Successfully downloaded file, size: {len(data)} bytes")
+            
+            print("Processing as Python file...")
+            answer = run_python(data.decode())
+            
+            print(f"Generated answer: {answer}")
+            state["final_answer"] = answer
+            state["current_step"] = "done"
+            return state
+        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
 
-        # For text-only questions, use the LLM
+    def _process_text(self, state: AgentState) -> AgentState:
+        """Process text-only questions using LLM."""
         print("\nProcessing as text-only question...")
         prompt = f"""
 Answer this question using the materials provided.
@@ -342,11 +376,77 @@ QUESTION:
         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"
+
     def _build_workflow(self) -> Graph:
+        """Build the workflow graph with conditional edges."""
         sg = StateGraph(state_schema=AgentState)
-        sg.add_node("answer", self._generate_answer)
-        sg.set_entry_point("answer")
-        sg.set_finish_point("answer")
+        
+        # Add nodes for each tool
+        sg.add_node("route", self._route_to_tool)
+        sg.add_node("process_image", self._process_image)
+        sg.add_node("process_video", self._process_video)
+        sg.add_node("process_spreadsheet", self._process_spreadsheet)
+        sg.add_node("process_python", self._process_python)
+        sg.add_node("process_text", self._process_text)
+        
+        # Set entry point
+        sg.set_entry_point("route")
+        
+        # Add conditional edges
+        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"
+            }
+        )
+        
+        # Set finish points for all tool nodes
+        for node in ["process_image", "process_video", "process_spreadsheet", "process_python", "process_text"]:
+            sg.set_finish_point(node)
+            
         return sg.compile()
 
 # ----------------------------------------------------------------------------------