Update inference.py

inference.py

@@ -11,7 +11,7 @@ API_BASE_URL = os.getenv("API_BASE_URL", "https://generativelanguage.googleapis.
 MODEL_NAME = os.getenv("MODEL_NAME", "gemini-2.5-flash") 
 
 # 2. Put your REAL Gemini key right here (starts with AIza...)
-GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "HF_TOKENS") 
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "KEYS_HERE") 
 
 # 3. Your live Hugging Face Space URL
 ENV_URL = os.getenv("ENV_URL", "https://rohannk-datacenter-openenv.hf.space")
@@ -30,9 +30,15 @@ def run_inference():
     except Exception as e:
         print(f"Error connecting to environment: {e}")
         return
-    
+
+    task_name = "DatacenterCooling"
+    print(f"[START] task={task_name}", flush=True)
+
     done = False
+    step = 0
+    total_current_score = 0.0
     while not done:
+        step += 1
         # Prepare prompt with current state
         prompt = f"""
         You are an AI managing a Data Center Cooling System.
@@ -62,7 +68,19 @@ def run_inference():
             
             state_resp = step_resp['state']
             done = step_resp['done']
-            print(f"Current Scores: {step_resp['scores']}\n")
+            
+            scores = step_resp.get('scores', {})
+            if isinstance(scores, dict):
+                current_reward = sum(scores.values())
+            elif isinstance(scores, (int, float)):
+                current_reward = float(scores)
+            else:
+                current_reward = 0.0
+                
+            total_current_score += current_reward
+            
+            print(f"[STEP] step={step} reward={current_reward}", flush=True)
+            print(f"Current Scores: {scores}\n")
             
             # Pause to prevent rate limits
             time.sleep(8) 
@@ -71,5 +89,7 @@ def run_inference():
             print(f"API Error: {e}")
             break
 
+    print(f"[END] task={task_name} score={total_current_score} steps={step}", flush=True)
+
 if __name__ == "__main__":
     run_inference()