Update app.py

app.py

@@ -54,27 +54,23 @@ def tighten(q: str) -> str:
 # Multimodal helpers
 # -------------------------
 
-def image_qa(image_path: str, prompt: str) -> str:
-    """Query MiniGPT-4-V for image-based QA."""
-    with open(image_path, "rb") as f:
-        data = {"prompt": prompt, "image": f.read()}
-        headers = {"Content-Type": "application/octet-stream"}
+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)
 
-def video_label(video_path: str, topk: int = 1) -> str:
-    """Get video classification using VideoMAE-Base-Short."""
-    with open(video_path, "rb") as f:
-        headers = {"Content-Type": "application/octet-stream"}
-        preds = client.post(
-            "MCG-NJU/videomae-base-short-finetuned-ucf101", 
-            data=f.read(),
-            headers=headers
-        )
-    preds = sorted(preds, key=lambda x: x["score"], reverse=True)[:topk]
-    return preds[0]["label"]
+def video_label_bytes(data: bytes) -> str:
+    """Get video classification using VideoMAE-Base-Short from bytes."""
+    headers = {"Content-Type": "application/octet-stream"}
+    preds = client.post(
+        "MCG-NJU/videomae-base-short-finetuned-ucf101", 
+        data=data,
+        headers=headers
+    )
+    return sorted(preds, key=lambda x: x["score"], reverse=True)[0]["label"]
 
-def sheet_answer(data: bytes, question: str) -> str:
-    """Process spreadsheet data and answer questions."""
+def sheet_answer_bytes(data: bytes, question: str) -> str:
+    """Process spreadsheet data from bytes and answer questions."""
     if mimetypes.guess_type("x.xlsx")[0] == "text/csv" or question.endswith(".csv"):
         df = pd.read_csv(io.BytesIO(data))
     else:
@@ -175,16 +171,14 @@ class BasicAgent:
                 # Download the file
                 file_data = self._download_file(file_url)
                 # Store in state
-                state["attachment_data"] = {
-                    "content": file_data,
-                    "type": self._detect_file_type(file_data, file_url)
-                }
-                # Set appropriate step based on file type
-                if state["attachment_data"]["type"] == "video":
+                state["attachment_data"] = file_data
+                # Detect type and set appropriate step
+                file_type = self._detect_file_type(file_url)
+                if file_type == "video":
                     state["current_step"] = "video"
-                elif state["attachment_data"]["type"] == "image":
+                elif file_type == "image":
                     state["current_step"] = "image"
-                elif state["attachment_data"]["type"] in ["excel", "csv"]:
+                elif file_type in ["excel", "csv"]:
                     state["current_step"] = "sheet"
                 return state
         except (json.JSONDecodeError, KeyError):
@@ -257,22 +251,23 @@ class BasicAgent:
             
         return state
 
-    def _detect_file_type(self, data: bytes, url: str) -> str:
+    def _detect_file_type(self, url: str) -> str:
         """Detect file type from URL extension."""
         ext = url.split(".")[-1].lower()
-        if ext in {"mp4"}: return "video"
-        if ext in {"jpg", "jpeg", "png"}: return "image"
-        if ext in {"xlsx"}: return "excel"
-        if ext in {"csv"}: return "csv"
-        return "unknown"
+        return {
+            "mp4": "video",
+            "jpg": "image",
+            "jpeg": "image",
+            "png": "image",
+            "xlsx": "excel",
+            "csv": "csv"
+        }.get(ext, "unknown")
 
     def _image_node(self, state: AgentState) -> AgentState:
         """Handle image-based questions."""
         try:
-            if "attachment_data" in state and "content" in state["attachment_data"]:
-                # Use the downloaded image data
-                image_data = state["attachment_data"]["content"]
-                answer = image_qa(image_data, "What is shown in this image?")
+            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")
@@ -285,10 +280,8 @@ class BasicAgent:
     def _video_node(self, state: AgentState) -> AgentState:
         """Handle video-based questions."""
         try:
-            if "attachment_data" in state and "content" in state["attachment_data"]:
-                # Use the downloaded video data
-                video_data = state["attachment_data"]["content"]
-                label = video_label(video_data)
+            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")
@@ -301,10 +294,8 @@ class BasicAgent:
     def _sheet_node(self, state: AgentState) -> AgentState:
         """Handle spreadsheet-based questions."""
         try:
-            if "attachment_data" in state and "content" in state["attachment_data"]:
-                # Use the downloaded spreadsheet data
-                sheet_data = state["attachment_data"]["content"]
-                answer = sheet_answer(sheet_data, state["question"])
+            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")
@@ -357,17 +348,9 @@ class BasicAgent:
     def _generate_answer(self, state: AgentState) -> AgentState:
         # Collect all relevant tool outputs
         materials = []
-        for step in state["history"]:
-            if step["step"] in {"search", "image", "video", "sheet", "code_analysis"}:
-                # Handle different output formats
-                if step["step"] == "search":
-                    output = step.get("results", [])
-                    if isinstance(output, list):
-                        output = "\n".join(output)
-                else:
-                    output = step.get("output", "")
-                # Format the output as JSON for better readability
-                materials.append(json.dumps(output, indent=2))
+        for h in state["history"]:
+            if h["step"] in {"search", "video", "image", "sheet", "code_analysis"}:
+                materials.append(json.dumps(h.get("output") or h.get("results"), indent=2))
         
         # Join all materials with proper formatting
         search_block = "\n".join(materials) if materials else "No artefacts available."