fix phase2 timeout

inference.py

@@ -1,132 +1,124 @@
+import os
 import threading
 import time
-import os
 import random
 import traceback
 from openai import OpenAI
 
 random.seed(42)
 
-# Start server
-try:
-    from env.fake_server import run_server
-    threading.Thread(target=run_server, daemon=True).start()
-    time.sleep(3)
-except Exception:
-    pass
-
+from env.fake_server import run_server
 from env.env import DeceptionEnv
 from env.attacker import simulate_attack
 
-API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
-MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-7B-Instruct")
-HF_TOKEN = os.getenv("HF_TOKEN")
 
-try:
+# If running inside HuggingFace Space → start server only
+if os.getenv("SPACE_ID"):
 
-    client = OpenAI(
-        base_url=API_BASE_URL,
-        api_key=HF_TOKEN
-    )
+    run_server()
 
-    env = DeceptionEnv()
 
-    # reset API must work
+else:
+
+    API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
+    MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-7B-Instruct")
+    HF_TOKEN = os.getenv("HF_TOKEN")
+
     try:
-        state = env.reset()
-    except Exception:
-        state = {}
 
-    print(
-        "[START] task=ai-deception env=cyber-security model=AI-agent",
-        flush=True
-    )
+        threading.Thread(target=run_server, daemon=True).start()
+        time.sleep(3)
+
+        client = OpenAI(
+            base_url=API_BASE_URL,
+            api_key=HF_TOKEN
+        )
+
+        env = DeceptionEnv()
+        state = env.reset()
 
-    rewards = []
-    history = []
+        print(
+            "[START] task=ai-deception env=cyber-security model=AI-agent",
+            flush=True
+        )
 
-    # EXACTLY 3 STEPS
-    for step in range(1, 4):
+        rewards = []
+        history = []
 
-        try:
-            simulate_attack()
-        except Exception:
-            pass
+        for step in range(1, 4):
 
-        # state API must work
-        try:
-            state = env.state()
-        except Exception:
-            pass
+            try:
+                simulate_attack()
+            except:
+                pass
 
-        summary = {
-            "failed_logins": state.get("failed_logins", 0),
-            "port_scans": state.get("port_scans", 0),
-            "suspicious_ips": len(state.get("suspicious_ips", []))
-        }
+            try:
+                state = env.state()
+            except:
+                pass
 
-        prompt = f"""
-You are a cybersecurity decision system.
+            summary = {
+                "failed_logins": state.get("failed_logins", 0),
+                "port_scans": state.get("port_scans", 0),
+                "suspicious_ips": len(state.get("suspicious_ips", []))
+            }
 
-Previous actions:
-{history}
+            prompt = f"""
+You are cybersecurity system.
 
-Current summary:
-{summary}
+Previous: {history}
+Current: {summary}
 
-Return ONLY one word:
+Return one:
 detect_attack
 deploy_honeypot
 block_ip
 """
 
-        try:
-            response = client.chat.completions.create(
-                model=MODEL_NAME,
-                messages=[{"role": "user", "content": prompt}],
-                temperature=0.2,
-                max_tokens=20,
-                timeout=15
-            )
+            try:
+                response = client.chat.completions.create(
+                    model=MODEL_NAME,
+                    messages=[{"role": "user", "content": prompt}],
+                    temperature=0.2,
+                    max_tokens=20,
+                    timeout=15
+                )
+
+                action = response.choices[0].message.content.strip()
 
-            action = response.choices[0].message.content.strip()
+            except:
 
-        except Exception:
-            # fallback logic
-            if not history:
-                action = "detect_attack"
-            elif history[-1] == "detect_attack":
-                action = "deploy_honeypot"
-            else:
-                action = "block_ip"
+                if not history:
+                    action = "detect_attack"
+                elif history[-1] == "detect_attack":
+                    action = "deploy_honeypot"
+                else:
+                    action = "block_ip"
 
-        history.append(action)
+            history.append(action)
 
-        try:
             state, reward, done, _ = env.step(action)
-        except Exception:
-            reward = 0.0
-            done = False
+            rewards.append(reward)
 
-        rewards.append(reward)
+            print(
+                f"[STEP] step={step} action={action} reward={reward:.2f} "
+                f"done={str(done).lower()} error=null",
+                flush=True
+            )
+
+        score = min(sum(rewards), 1.0)
 
         print(
-            f"[STEP] step={step} action={action} reward={reward:.2f} "
-            f"done={str(done).lower()} error=null",
+            f"[END] success=true steps={len(rewards)} "
+            f"score={score:.2f} rewards={','.join(f'{r:.2f}' for r in rewards)}",
             flush=True
         )
 
-    score = min(sum(rewards), 1.0)
+    except Exception:
 
-    print(
-        f"[END] success=true steps={len(rewards)} "
-        f"score={score:.2f} rewards={','.join(f'{r:.2f}' for r in rewards)}",
-        flush=True
-    )
+        traceback.print_exc()
 
-except Exception:
-    traceback.print_exc()
-    print(
-        "[END] success=false steps=0 score=0.00 rewards=",
-        flush=True
-    )
+        print(
+            "[END] success=false steps=0 score=0.00 rewards=",
+            flush=True
+        )