Update app.py

app.py

@@ -1,7 +1,6 @@
 import os
 import re
 import base64
-import json
 from io import BytesIO
 from pathlib import Path
 from zipfile import ZipFile
@@ -65,6 +64,65 @@ def filename_from_response(response: requests.Response, task_id: str) -> str:
     return f"{clean_filename(task_id)}{extension}"
 
 
+def download_gaia_attachment(task_id: str) -> tuple[bytes, str]:
+    """Baixa um anexo pela API do curso, com fallback para o dataset oficial."""
+    task_id = str(task_id).strip()
+    course_url = f"{DEFAULT_API_URL}/files/{task_id}"
+    response = requests.get(course_url, timeout=HTTP_TIMEOUT)
+    if response.ok:
+        return response.content, filename_from_response(response, task_id)
+    if response.status_code != 404:
+        response.raise_for_status()
+
+    questions_response = requests.get(
+        f"{DEFAULT_API_URL}/questions", timeout=HTTP_TIMEOUT
+    )
+    questions_response.raise_for_status()
+    item = next(
+        (
+            question
+            for question in questions_response.json()
+            if str(question.get("task_id")) == task_id
+        ),
+        None,
+    )
+    filename = str((item or {}).get("file_name") or "").strip()
+    if not filename:
+        raise FileNotFoundError(f"No attachment exists for task {task_id}.")
+
+    token = os.getenv("HF_TOKEN")
+    if not token:
+        raise RuntimeError(
+            "The course file endpoint returned 404 and HF_TOKEN is required "
+            "for the official GAIA dataset fallback."
+        )
+
+    from huggingface_hub import hf_hub_download
+
+    errors = []
+    for dataset_path in (
+        f"2023/validation/{filename}",
+        filename,
+    ):
+        try:
+            local_path = hf_hub_download(
+                repo_id="gaia-benchmark/GAIA",
+                repo_type="dataset",
+                filename=dataset_path,
+                token=token,
+            )
+            return Path(local_path).read_bytes(), filename
+        except Exception as exc:
+            errors.append(str(exc))
+
+    raise RuntimeError(
+        "Could not download the attachment from the course API or official "
+        "GAIA dataset. Accept the dataset access conditions at "
+        "https://huggingface.co/datasets/gaia-benchmark/GAIA and ensure "
+        f"HF_TOKEN has read access. Details: {' | '.join(errors)}"
+    )
+
+
 def extract_attachment_text(data: bytes, filename: str) -> str:
     """Extrai conteúdo legível dos formatos mais comuns do GAIA."""
     suffix = Path(filename).suffix.lower()
@@ -180,14 +238,9 @@ class InspectGaiaAttachmentTool(Tool):
         if not task_id:
             return "No task_id was supplied."
 
-        url = f"{DEFAULT_API_URL}/files/{task_id}"
         try:
-            response = requests.get(url, timeout=HTTP_TIMEOUT)
-            if response.status_code == 404:
-                return f"No attachment exists for task {task_id}."
-            response.raise_for_status()
-            filename = filename_from_response(response, task_id)
-            extracted = extract_attachment_text(response.content, filename)
+            data, filename = download_gaia_attachment(task_id)
+            extracted = extract_attachment_text(data, filename)
             return f"Attachment filename: {filename}\n\n{extracted}"
         except Exception as exc:
             return f"Could not inspect attachment for task {task_id}: {exc}"
@@ -262,13 +315,16 @@ class AnalyzeGaiaImageTool(Tool):
             return "Image analysis failed: HF_TOKEN is not configured."
 
         try:
-            response = requests.get(
-                f"{DEFAULT_API_URL}/files/{str(task_id).strip()}",
-                timeout=HTTP_TIMEOUT,
-            )
-            response.raise_for_status()
-            mime = response.headers.get("content-type", "image/png").split(";")[0]
-            encoded = base64.b64encode(response.content).decode("ascii")
+            data, filename = download_gaia_attachment(str(task_id).strip())
+            suffix = Path(filename).suffix.lower()
+            mime = {
+                ".png": "image/png",
+                ".jpg": "image/jpeg",
+                ".jpeg": "image/jpeg",
+                ".webp": "image/webp",
+                ".gif": "image/gif",
+            }.get(suffix, "image/png")
+            encoded = base64.b64encode(data).decode("ascii")
 
             client = OpenAI(
                 base_url="https://router.huggingface.co/v1",
@@ -624,46 +680,66 @@ Formatting rules:
   ordering, capitalization, and plurality.
 - If asked for a quote: return only the requested spoken words.
 - If asked for a chess move: return only algebraic notation.
-- Never include labels, Markdown, citations, rationale, or "FINAL ANSWER".
-
-Return one JSON object and nothing else:
-{{
-  "final_answer": "exact value to submit",
-  "changed": true,
-  "review_note": "brief reason, maximum 20 words"
-}}
+- Never include labels, Markdown, citations, rationale, or "FINAL ANSWER" inside
+  the final_answer value.
+
+Return exactly these two XML-style fields:
+<final_answer>exact value to submit</final_answer>
+<review_note>brief reason, maximum 20 words</review_note>
+
+Do not return JSON. Do not wrap the fields in Markdown or a code block.
 """.strip()
 
         last_error = None
-        for use_json_mode in (True, False):
+        for attempt in range(2):
             try:
-                kwargs = {
-                    "model": reviewer_model,
-                    "api_key": self.gemini_api_key,
-                    "messages": [{"role": "user", "content": review_prompt}],
-                    "temperature": 0,
-                    "max_tokens": 350,
-                }
-                if use_json_mode:
-                    kwargs["response_format"] = {"type": "json_object"}
-                response = completion(**kwargs)
+                retry_instruction = (
+                    ""
+                    if attempt == 0
+                    else "\nIMPORTANT: Return both XML fields exactly as specified."
+                )
+                response = completion(
+                    model=reviewer_model,
+                    api_key=self.gemini_api_key,
+                    messages=[
+                        {
+                            "role": "user",
+                            "content": review_prompt + retry_instruction,
+                        }
+                    ],
+                    temperature=0,
+                    max_tokens=350,
+                )
                 content = str(response.choices[0].message.content).strip()
                 content = re.sub(
-                    r"^```(?:json)?\s*|\s*```$", "", content, flags=re.I
+                    r"^```(?:xml|text)?\s*|\s*```$", "", content, flags=re.I
+                )
+                answer_match = re.search(
+                    r"<final_answer>\s*(.*?)\s*</final_answer>",
+                    content,
+                    flags=re.I | re.S,
                 )
-                start = content.find("{")
-                end = content.rfind("}")
-                if start >= 0 and end > start:
-                    content = content[start : end + 1]
-                payload = json.loads(content)
+                note_match = re.search(
+                    r"<review_note>\s*(.*?)\s*</review_note>",
+                    content,
+                    flags=re.I | re.S,
+                )
+                if not answer_match:
+                    raise ValueError(
+                        "Gemini did not return the <final_answer> field."
+                    )
                 final_answer = self.deterministic_answer_cleanup(
-                    str(payload.get("final_answer", ""))
+                    answer_match.group(1)
                 )
                 if not final_answer:
                     raise ValueError("Gemini returned an empty final_answer.")
 
                 changed = final_answer != candidate
-                note = str(payload.get("review_note", "")).strip()
+                note = (
+                    note_match.group(1).strip()
+                    if note_match
+                    else "Review completed without a note."
+                )
                 print(f"Candidate answer: {candidate}")
                 print(f"Gemini reviewed answer: {final_answer}")
                 print(f"Gemini changed answer: {changed}")