Create inference.py

inference.py

@@ -0,0 +1,91 @@
+import os
+import json
+from openai import OpenAI
+from src.disaster_grid.environment import CityGrid
+from src.disaster_grid.models import AgentAction, ActionType
+
+def run_inference():
+    print("Initializing Disaster Grid Environment...")
+    env = CityGrid()
+    
+    # Initialize the API client
+    # NOTE: If the hackathon requires a specific API (like Grok or TogetherAI), 
+    # just change the base_url and model name below!
+    client = OpenAI(
+        api_key=os.environ.get("API_KEY", "your-api-key-here"),
+        base_url="https://api.openai.com/v1" 
+    )
+    
+    print("\n--- Starting Disaster Scenario ---")
+    # Our environment returns a tuple: (observation, info) on reset
+    obs, _ = env.reset()
+    
+    done = False
+    
+    while not done:
+        print(f"\nTime Step: {env.step_count}/50 | Energy: {env.agent_energy}")
+        
+        # 1. Package the environment state into a prompt for the LLM
+        prompt = f"""
+        You are an Autonomous AI Emergency Manager.
+        
+        Current Environment State:
+        {json.dumps(obs, indent=2)}
+        
+        Rules:
+        - You are on a 5x5 grid (indices 0 to 24). You start at index 0.
+        - Moving (MOVE_N, MOVE_S, MOVE_E, MOVE_W) costs 2 energy.
+        - REPAIR costs 15 energy and adds 25 health to your current sector.
+        - RECHARGE adds 20 energy, but ONLY works if you are at index 0 (Base).
+        - Do not let your energy hit 0. Navigate to critical sectors and repair them.
+        
+        Determine the best action. You MUST respond with a perfectly formatted JSON object matching this schema:
+        {{"action": "MOVE_N" | "MOVE_S" | "MOVE_E" | "MOVE_W" | "REPAIR" | "RECHARGE" | "WAIT", "reasoning": "<string explaining your strategy>"}}
+        """
+        
+        try:
+            # 2. Call the LLM
+            response = client.chat.completions.create(
+                model="gpt-4o", # Replace with "grok-beta" or your required model
+                messages=[
+                    {"role": "system", "content": "You are a JSON-only API. You only output raw, valid JSON."},
+                    {"role": "user", "content": prompt}
+                ],
+                response_format={"type": "json_object"}
+            )
+            
+            # 3. Parse the JSON response
+            raw_response = response.choices[0].message.content
+            action_data = json.loads(raw_response)
+            
+            # Validate it through our Pydantic model just to be safe
+            action_parsed = AgentAction(**action_data)
+            
+            print(f"🤖 AI decided: {action_parsed.action.value}")
+            print(f"   Reasoning: {action_parsed.reasoning}")
+            
+            # 4. Execute the action in the environment
+            # Our env.step returns a 5-item tuple and handles the dict parsing internally
+            obs, reward, done, truncated, info = env.step(action_data)
+            
+            # Print any errors from the environment engine (like wall bumps)
+            step_result = info.get("step_result", {})
+            if step_result.get("is_error"):
+                print(f"⚠️  Engine Warning: {step_result.get('error_message')}")
+            
+        except Exception as e:
+            print(f"❌ Error during LLM processing: {e}")
+            print("Forcing a WAIT action to prevent the loop from crashing...")
+            fallback_action = {"action": ActionType.WAIT.value, "reasoning": "Fallback due to error"}
+            obs, reward, done, truncated, info = env.step(fallback_action)
+
+    # 5. The episode is finished. Print the final summary!
+    print("\n" + "="*40)
+    print("🎉 EPISODE COMPLETE 🎉")
+    print(f"Final City Health: {sum(env.grid_health)/25:.1f}/100")
+    print(f"Final Energy:      {env.agent_energy}")
+    print(f"Steps Taken:       {env.step_count}")
+    print("="*40)
+
+if __name__ == "__main__":
+    run_inference()