Update agent.py

agent.py

@@ -2,45 +2,66 @@ import os
 import requests
 from typing import Dict, List
 
-from run import run_and_submit_all
+from run import run_and_submit_all  # Adjust path if needed
 
 class GaiaAgent:
     def __init__(self):
-        self.api_url = os.environ["HF_MISTRAL_ENDPOINT"]
-        self.api_key = os.environ["HF_TOKEN"]
+        self.api_url = os.environ.get("HF_MISTRAL_ENDPOINT")
+        self.api_key = os.environ.get("HF_TOKEN")
+        self.model_id = os.environ.get("LLM_MODEL_ID")
+
+        assert self.api_url, "❌ HF_MISTRAL_ENDPOINT is missing!"
+        assert self.api_key, "❌ HF_TOKEN is missing!"
+        assert self.model_id, "❌ LLM_MODEL_ID is missing!"
+
+        print(f"✅ [INIT] Model ID: {self.model_id}")
+        print(f"✅ [INIT] Endpoint: {self.api_url}")
+
         self.headers = {
             "Authorization": f"Bearer {self.api_key}",
             "Content-Type": "application/json",
         }
 
     def generate(self, prompt: str, stop: List[str] = []) -> str:
-        print("📤 [generate()] Sending prompt to model.")
-        print(prompt[:200])  # Print first 200 chars of the prompt for sanity
-        
+        print("🧠 [GENERATE] Prompt sent to model:")
+        print(prompt)
+
         payload = {
             "inputs": prompt,
             "parameters": {
-            "temperature": 0.0,
-            "max_new_tokens": 1024,
-            "stop": stop,
+                "temperature": 0.0,
+                "max_new_tokens": 1024,
+                "stop": stop,
+            }
         }
-    }
-        response = requests.post(self.api_url, headers=self.headers, json=payload)
-        response.raise_for_status()
+
+        try:
+            response = requests.post(self.api_url, headers=self.headers, json=payload)
+            response.raise_for_status()
+        except Exception as e:
+            print(f"❌ [ERROR] Request failed: {e}")
+            return "ERROR: Model call failed"
+
         output = response.json()
-        
+        print(f"✅ [RESPONSE] Raw output: {output}")
+
         if isinstance(output, dict) and "generated_text" in output:
             return output["generated_text"]
         elif isinstance(output, list) and "generated_text" in output[0]:
             return output[0]["generated_text"]
-        return str(output)
+        else:
+            return str(output)
 
     def answer_question(self, question: Dict) -> str:
         q = question["question"]
+        print(f"📌 [QUESTION] ID: {question.get('id', 'N/A')} - {q}")
+
         prompt = f"""You are a helpful agent answering a science question.
 Question: {q}
 Answer:"""
+
         return self.generate(prompt).strip()
 
     def run(self):
+        print("🚀 [RUN] Starting submission...")
         run_and_submit_all(self.answer_question)