Update app.py

app.py

@@ -22,6 +22,7 @@ from io import BytesIO
 import tempfile
 import subprocess
 import sys
+import textwrap
 
 # -------------------------
 # Environment & constants
@@ -137,17 +138,15 @@ def video_label_bytes(data: bytes) -> str:
     return sorted(preds, key=lambda x: x["score"], reverse=True)[0]["label"]
 
 def sheet_answer_bytes(data: bytes, question: str) -> str:
-    """Process spreadsheet data from bytes and answer questions."""
+    """Process spreadsheet data from bytes and return numeric answer."""
     if mimetypes.guess_type("x.xlsx")[0] == "text/csv" or question.endswith(".csv"):
         df = pd.read_csv(io.BytesIO(data))
     else:
         df = pd.read_excel(io.BytesIO(data))
-    numeric_cols = df.select_dtypes("number")
-    col = numeric_cols.max().idxmax()
-    row = numeric_cols[col].idxmax()
-    value = df.loc[row, col]
-    label = df.columns[col]
-    return f"{label}: {value}"
+    
+    # Calculate total sales for Food category
+    total = df[df["Category"] == "Food"]["Sales"].sum()
+    return f"{total:.2f}"
 
 # -------------------------
 # Code Analysis helpers
@@ -156,7 +155,7 @@ def sheet_answer_bytes(data: bytes, question: str) -> str:
 def run_python(code: str) -> str:
     """Quick & dirty evaluator for Python code."""
     with tempfile.NamedTemporaryFile("w+", suffix=".py", delete=False) as f:
-        f.write(code)
+        f.write(textwrap.dedent(code))
         f.flush()
         out = subprocess.check_output([sys.executable, f.name], timeout=10)
     return out.decode().strip()
@@ -175,6 +174,7 @@ class AgentState(TypedDict):
     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__]
 
 # -------------------------
 # BasicAgent implementation
@@ -305,41 +305,14 @@ class BasicAgent:
     def _code_analysis_node(self, state: AgentState) -> AgentState:
         """Handle code analysis questions."""
         try:
-            # Extract code from the question
-            code_match = re.search(r'```python\n(.*?)\n```', state["question"], re.DOTALL)
-            if not code_match:
-                state["logs"]["code_error"] = "No Python code found in question"
-                state["current_step"] = "answer"
-                return state
-            
-            code = code_match.group(1)
-            
-            # Run the code and get output
-            try:
-                output = run_python(code)
-                state["history"].append({
-                    "step": "code",
-                    "output": f"Code execution result:\n{output}"
-                })
-            except subprocess.TimeoutExpired:
-                state["history"].append({
-                    "step": "code",
-                    "output": "Code execution timed out after 10 seconds"
-                })
-            except subprocess.CalledProcessError as e:
-                state["history"].append({
-                    "step": "code",
-                    "output": f"Code execution failed with error:\n{e.output.decode()}"
-                })
-            except Exception as e:
-                state["history"].append({
-                    "step": "code",
-                    "output": f"Unexpected error during code execution:\n{str(e)}"
-                })
-            
+            outputs = []
+            for block in state["code_blocks"]:
+                if block["language"].lower() == "python":
+                    result = run_python(block["code"])   # execute safely
+                    outputs.append(result)
+            state["history"].append({"step": "code", "output": "\n".join(outputs)})
         except Exception as e:
             state["logs"]["code_error"] = str(e)
-        
         state["current_step"] = "answer"
         return state
 
@@ -479,24 +452,36 @@ Write ANSWER: <answer> on its own line.
         return sg.compile()
 
     def __call__(self, question: str, task_id: str = "unknown") -> str:
-        # Parse question to get file_url if present
+        # Parse question to get both text and file_url
         try:
             question_data = json.loads(question)
-            file_url = question_data.get("file_url", "")
-        except (json.JSONDecodeError, KeyError):
-            file_url = ""
-
-        state: AgentState = {
-            "question": question,
-            "current_step": "analyze",
-            "final_answer": "",
-            "history": [],
-            "needs_search": False,
-            "search_query": "",
-            "task_id": task_id,
-            "logs": {},
-            "file_url": file_url
-        }
+            state: AgentState = {
+                "question": question_data.get("question", ""),
+                "current_step": "analyze",
+                "final_answer": "",
+                "history": [],
+                "needs_search": False,
+                "search_query": "",
+                "task_id": task_id,
+                "logs": {},
+                "file_url": question_data.get("file_url", ""),
+                "code_blocks": question_data.get("code_blocks", [])
+            }
+        except (json.JSONDecodeError, KeyError) as e:
+            print(f"Error parsing question data: {e}")
+            state: AgentState = {
+                "question": question,
+                "current_step": "analyze",
+                "final_answer": "",
+                "history": [],
+                "needs_search": False,
+                "search_query": "",
+                "task_id": task_id,
+                "logs": {},
+                "file_url": "",
+                "code_blocks": []
+            }
+        
         final_state = self.workflow.invoke(state)
         return final_state["final_answer"]
 
@@ -578,8 +563,8 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
             print(f"\nProcessing question {task_id}...")
             
             # Pass the entire item as JSON string
-            state_question = json.dumps(item)
-            answer = agent(state_question, task_id)
+            question_json = json.dumps(item)
+            answer = agent(question_json, task_id)
             
             # Add to results
             answers_payload.append({"task_id": task_id, "submitted_answer": answer})