Update app.py

app.py

@@ -53,11 +53,30 @@ def tighten(q: str) -> str:
 # Multimodal helpers
 # -------------------------
 
+def retry_hf_inference(func):
+    """Decorator to retry HF Inference API calls with backoff."""
+    def wrapper(*args, **kwargs):
+        max_retries = 2
+        base_delay = 7
+        
+        for attempt in range(max_retries + 1):
+            try:
+                return func(*args, **kwargs)
+            except Exception as e:
+                if attempt == max_retries:
+                    raise
+                delay = base_delay * (attempt + 1)
+                print(f"HF API error: {str(e)}. Retrying in {delay}s...")
+                time.sleep(delay)
+    return wrapper
+
+@retry_hf_inference
 def image_qa_bytes(data: bytes, prompt: str) -> str:
     """Query MiniGPT-4-V for image-based QA using bytes."""
     headers = {"Content-Type": "application/octet-stream"}
     return client.post("Vision-CAIR/MiniGPT4-V", data=data, headers=headers)
 
+@retry_hf_inference
 def video_label_bytes(data: bytes) -> str:
     """Get video classification using VideoMAE-Base-Short from bytes."""
     headers = {"Content-Type": "application/octet-stream"}
@@ -94,7 +113,7 @@ class AgentState(TypedDict):
     search_query: Annotated[str, override]
     task_id: Annotated[str, override]
     logs: Annotated[Dict[str, Any], merge_dicts]
-    attachment_data: Annotated[Dict[str, bytes], merge_dicts]
+    file_url: Annotated[str, override]
 
 # -------------------------
 # BasicAgent implementation
@@ -128,10 +147,8 @@ class BasicAgent:
         try:
             question_data = json.loads(state["question"])
             if "file_url" in question_data:
-                file_url = question_data["file_url"]
-                file_data = self._download_file(file_url)
-                state["attachment_data"] = file_data
-                file_type = self._detect_file_type(file_url)
+                state["file_url"] = question_data["file_url"]
+                file_type = self._detect_file_type(state["file_url"])
                 if file_type == "video":
                     state["current_step"] = "video"
                 elif file_type == "image":
@@ -180,43 +197,34 @@ class BasicAgent:
     def _image_node(self, state: AgentState) -> AgentState:
         """Handle image-based questions."""
         try:
-            if "attachment_data" in state:
-                answer = image_qa_bytes(state["attachment_data"], "What is shown in this image?")
-                state["history"].append({"step": "image", "output": answer})
-            else:
-                raise ValueError("No image data found in state")
-            state["current_step"] = "answer"
+            data = self._download_file(state["file_url"])
+            answer = image_qa_bytes(data, "What is shown in this image?")
+            state["history"].append({"step": "image", "output": answer})
         except Exception as e:
             state["logs"]["image_error"] = str(e)
-            state["current_step"] = "answer"
+        state["current_step"] = "answer"
         return state
 
     def _video_node(self, state: AgentState) -> AgentState:
         """Handle video-based questions."""
         try:
-            if "attachment_data" in state:
-                label = video_label_bytes(state["attachment_data"])
-                state["history"].append({"step": "video", "output": label})
-            else:
-                raise ValueError("No video data found in state")
-            state["current_step"] = "answer"
+            data = self._download_file(state["file_url"])
+            label = video_label_bytes(data)
+            state["history"].append({"step": "video", "output": label})
         except Exception as e:
             state["logs"]["video_error"] = str(e)
-            state["current_step"] = "answer"
+        state["current_step"] = "answer"
         return state
 
     def _sheet_node(self, state: AgentState) -> AgentState:
         """Handle spreadsheet-based questions."""
         try:
-            if "attachment_data" in state:
-                answer = sheet_answer_bytes(state["attachment_data"], state["question"])
-                state["history"].append({"step": "sheet", "output": answer})
-            else:
-                raise ValueError("No spreadsheet data found in state")
-            state["current_step"] = "answer"
+            data = self._download_file(state["file_url"])
+            answer = sheet_answer_bytes(data, state["file_url"])
+            state["history"].append({"step": "sheet", "output": answer})
         except Exception as e:
             state["logs"]["sheet_error"] = str(e)
-            state["current_step"] = "answer"
+        state["current_step"] = "answer"
         return state
 
     def _perform_search(self, state: AgentState) -> AgentState:
@@ -243,13 +251,24 @@ class BasicAgent:
         return state
 
     def _generate_answer(self, state: AgentState) -> AgentState:
-        # Collect all relevant tool outputs
+        # Collect all tool outputs with clear section headers
         materials = []
         for h in state["history"]:
-            if h["step"] in {"search", "video", "image", "sheet"}:
-                materials.append(json.dumps(h.get("output") or h.get("results"), indent=2))
+            if h["step"] == "search":
+                materials.append("=== Search Results ===")
+                for result in h.get("results", []):
+                    materials.append(result)
+            elif h["step"] == "image":
+                materials.append("=== Image Analysis ===")
+                materials.append(h.get("output", ""))
+            elif h["step"] == "video":
+                materials.append("=== Video Analysis ===")
+                materials.append(h.get("output", ""))
+            elif h["step"] == "sheet":
+                materials.append("=== Spreadsheet Analysis ===")
+                materials.append(h.get("output", ""))
         
-        search_block = "\n".join(materials) if materials else "No artefacts available."
+        search_block = "\n\n".join(materials) if materials else "No materials available."
         
         prompt = f"""
 Answer this question using ONLY the materials provided.
@@ -325,7 +344,7 @@ Write ANSWER: <answer> on its own line.
             "search_query": "",
             "task_id": task_id,
             "logs": {},
-            "attachment_data": {}
+            "file_url": ""
         }
         final_state = self.workflow.invoke(state)
         return final_state["final_answer"]
@@ -418,7 +437,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
                 "search_query": "",
                 "task_id": task_id,
                 "logs": {},
-                "attachment_data": {}
+                "file_url": ""
             }
             
             # Run the workflow