Update app.py

app.py

@@ -6,7 +6,7 @@ import json
 import time
 import pandas as pd
 from datetime import datetime
-from typing import List, Dict, Any, Annotated
+from typing import List, Dict, Any, Annotated, Optional
 from langgraph.graph import Graph, StateGraph
 from typing_extensions import TypedDict
 from openai import OpenAI
@@ -35,6 +35,9 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 # Initialize HF client
 client = InferenceClient(token=HF_TOKEN)
 
+# Create a single Session for all requests
+SESSION = requests.Session()
+
 # -------------------------
 # Constants
 # -------------------------
@@ -143,20 +146,17 @@ 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:
+def sheet_answer_bytes(data: bytes) -> str:
     """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:
+    try:
         df = pd.read_excel(io.BytesIO(data))
+    except ValueError:
+        df = pd.read_csv(io.BytesIO(data))
     
-    # Calculate total sales for Food category
-    total = df[df["Category"] == "Food"]["Sales"].sum()
-    return f"{total:.2f}"
-
-# -------------------------
-# Code Analysis helpers
-# -------------------------
+    if {"Category", "Sales"}.issubset(df.columns):
+        total = df[df["Category"] == "Food"]["Sales"].sum()
+        return f"{total:.2f}"
+    return "sheet_answer_placeholder"
 
 def run_python(code: str) -> str:
     """Quick & dirty evaluator for Python code."""
@@ -166,6 +166,17 @@ def run_python(code: str) -> str:
         out = subprocess.check_output([sys.executable, f.name], timeout=10)
     return out.decode().strip()
 
+def discover_attachment(task_id: str, api_url: str) -> Optional[str]:
+    """Probe if a task has an attachment, return URL if it exists."""
+    probe = f"{api_url}/files/{task_id}"
+    try:
+        r = SESSION.get(probe, stream=True, timeout=10, allow_redirects=True)
+        if 200 <= r.status_code < 400:
+            return probe
+    except requests.RequestException:
+        pass
+    return None
+
 # -------------------------
 # State definition
 # -------------------------
@@ -187,12 +198,16 @@ class AgentState(TypedDict):
 # -------------------------
 
 class BasicAgent:
-    def __init__(self, session: requests.Session):
+    """A very small agent that can handle text questions and a few file types."""
+
+    JSON_INSTRUCTION = "Return ONLY this exact JSON object: {\"ANSWER\": \"<answer text>\"}"
+
+    def __init__(self, api_url: str = DEFAULT_API_URL):
         if not OPENAI_API_KEY:
             raise EnvironmentError("OPENAI_API_KEY not set")
         self.llm = OpenAI(api_key=OPENAI_API_KEY)
+        self.api_url = api_url
         self.workflow = self._build_workflow()
-        self.session = session
 
     def _call_llm(self, prompt: str, max_tokens: int = 256) -> str:
         try:
@@ -211,20 +226,25 @@ class BasicAgent:
             raise
 
     def _safe_parse(self, raw: str) -> str:
+        """Pull ANSWER from the JSON string, tolerant to model chatter."""
         try:
             return json.loads(raw)["ANSWER"]
         except Exception:
-            # grab the first {...} in the text
+            # Try to find any JSON object in the text
             match = re.search(r'\{.*?\}', raw, re.S)
             if match:
                 try:
                     return json.loads(match.group())["ANSWER"]
                 except Exception:
                     pass
-            # as a last resort, strip everything before the first colon
+            # As a last resort, take everything after the first colon
             return raw.split(':', 1)[-1].strip()
 
     def __call__(self, question: str, task_id: str = "unknown", file_url: str = "") -> str:
+        # 1) if file_url blank, attempt discovery once
+        if not file_url:
+            file_url = discover_attachment(task_id, self.api_url) or ""
+
         state: AgentState = {
             "question": question,
             "current_step": "answer",
@@ -249,7 +269,7 @@ class BasicAgent:
         if state["file_url"]:
             try:
                 print(f"Downloading {state['file_url']} …")
-                response = self.session.get(state["file_url"], timeout=30)
+                response = SESSION.get(state["file_url"], timeout=30)
                 response.raise_for_status()
                 data = response.content
                 print(f"Successfully downloaded file, size: {len(data)} bytes")
@@ -277,7 +297,7 @@ class BasicAgent:
                     answer = video_label_bytes(data)
                 elif "spreadsheet" in kind or "excel" in kind:
                     print("Processing as spreadsheet...")
-                    answer = sheet_answer_bytes(data, state["question"])
+                    answer = sheet_answer_bytes(data)
                 elif state["file_url"].endswith(".py"):
                     print("Processing as Python file...")
                     answer = run_python(data.decode())
@@ -308,8 +328,7 @@ Answer this question using the materials provided.
 QUESTION:
 {state['question']}
 
-Return ONLY this exact JSON object:
-{{"ANSWER": "<answer text>"}}
+{self.JSON_INSTRUCTION}
 """
         try:
             raw = self._call_llm(prompt, 300)