Update agent.py

agent.py

@@ -14,12 +14,12 @@ from google import genai
 from google.genai import types as genai_types
 from smolagents import CodeAgent, DuckDuckGoSearchTool, OpenAIModel, Tool
 
-
 load_dotenv()
 
 DEFAULT_SCORING_API_URL = "https://agents-course-unit4-scoring.hf.space"
 DEFAULT_GEMINI_MODEL_ID = os.getenv("GEMINI_MODEL_ID", "gemini/gemini-2.5-flash")
 GEMINI_OPENAI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
+
 HTTP_SESSION = requests.Session()
 HTTP_SESSION.headers.update(
     {
@@ -27,8 +27,8 @@ HTTP_SESSION.headers.update(
         "Accept": "application/json, text/plain, */*",
     }
 )
-GAIA_OUTPUT_PROMPT = """
-You are an evaluation-grade GAIA benchmark AI assistant. Your execution must be perfectly precise and highly deterministic. Your sole purpose is to isolate and output the exact, minimal final answer.
+
+GAIA_OUTPUT_PROMPT = """You are an evaluation-grade GAIA benchmark AI assistant. Your execution must be perfectly precise and highly deterministic. Your sole purpose is to isolate and output the exact, minimal final answer.
 
 ### CORE DIRECTIVE
 You must NEVER output explanations, intermediate reasoning, conversational filler, or comments.
@@ -59,13 +59,10 @@ Format: [FINAL_ANSWER]
 ### EXAMPLES
 Q: What is 2 + 2?
 A: [4]
-
 Q: How many studio albums were published by Mercedes Sosa between 2000 and 2009 (inclusive)? Use 2022 English Wikipedia.
 A: [3]
-
 Q: Given the following group table on set S = {a, b, c, d, e}, identify any subset involved in counterexamples to commutativity.
 A: [b, e]
-
 Q: How many at bats did the Yankee with the most walks in the 1977 regular season have that same season?
 A: [519]
 """.strip()
@@ -95,7 +92,6 @@ def _extract_text(response) -> str:
     text = getattr(response, "text", None)
     if text:
         return text.strip()
-
     parts: list[str] = []
     for candidate in getattr(response, "candidates", []) or []:
         content = getattr(candidate, "content", None)
@@ -110,7 +106,6 @@ def _call_gemini_text(prompt: str, system_instruction: str | None = None) -> str
     config = genai_types.GenerateContentConfig(temperature=0)
     if system_instruction:
         config.system_instruction = system_instruction
-
     response = _genai_client().models.generate_content(
         model=_normalize_model_id(),
         contents=prompt,
@@ -157,18 +152,15 @@ def _query_uploaded_file(file_bytes: bytes, file_name: str, mime_type: str, user
     suffix = Path(file_name or "attachment").suffix
     temp_path = None
     uploaded_file = None
-
     try:
         with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
             temp_file.write(file_bytes)
             temp_path = temp_file.name
-
         uploaded_file = client.files.upload(
             file=temp_path,
             config={"mime_type": mime_type},
         )
         _wait_until_active(client, uploaded_file)
-
         response = client.models.generate_content(
             model=_normalize_model_id(),
             contents=[uploaded_file, user_query],
@@ -270,7 +262,6 @@ class MathSolver(Tool):
             result = eval_node(parsed)
         except Exception as exc:
             return f"Math error: {exc}"
-
         if isinstance(result, float) and result.is_integer():
             return str(int(result))
         return str(result)
@@ -466,26 +457,31 @@ class FileAttachmentQueryTool(Tool):
         file_bytes, content_type = _download_task_file(task_id)
         suffix = Path(file_name or "").suffix.lower()
         mime_type = _mime_from_name(file_name, content_type)
-
         if suffix in {".txt", ".md", ".json", ".csv", ".py", ".html", ".xml", ".yaml", ".yml", ".log"}:
             file_text = _decode_text_bytes(file_bytes)
             return _textual_file_answer(file_text, user_query)
-
         if suffix in {".xlsx", ".xls"}:
             file_text = _excel_to_text(file_bytes)
             return _textual_file_answer(file_text, user_query)
-
         return _query_uploaded_file(file_bytes, file_name, mime_type, user_query)
 
 
 class BasicAgent:
     def __init__(self):
+        # We keep this here to ensure Gemini Tools don't throw an error later
         _get_gemini_api_key()
+        
+        # Get the OpenAI API key for the main loop
+        openai_api_key = os.getenv("OPENAI_API_KEY")
+        if not openai_api_key:
+            raise RuntimeError(
+                "Missing OpenAI API key. Set OPENAI_API_KEY in your Hugging Face Space secrets."
+            )
+
         self.agent = CodeAgent(
             model=OpenAIModel(
-                model_id=_normalize_model_id(),
-                api_base=GEMINI_OPENAI_BASE_URL,
-                api_key=_get_gemini_api_key(),
+                model_id="gpt-4o-mini",
+                api_key=openai_api_key,
                 temperature=0,
             ),
             tools=[
@@ -522,7 +518,6 @@ class BasicAgent:
                 f"- file_name: {file_name}\n"
                 "Use run_query_with_file if the question requires the attachment."
             )
-
         result = self.agent.run("\n".join(prompt_parts).strip())
         return _normalize_final_answer(str(result))
 
@@ -532,4 +527,4 @@ if __name__ == "__main__":
         "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? "
         "You can use the latest 2022 version of english wikipedia."
     )
-    print(BasicAgent()(sample_question))
+    print(BasicAgent()(sample_question))