fix: update HF router config and add required logs

inference.py

@@ -20,20 +20,22 @@ logger = logging.getLogger(__name__)
 
 def setup_openai_client() -> tuple[OpenAI, str]:
     """Initialize OpenAI client from environment variables."""
-    api_base_url = os.getenv('API_BASE_URL')
-    model_name = os.getenv('MODEL_NAME', 'gpt-3.5-turbo')
+    api_base_url = os.getenv('API_BASE_URL', 'https://router.huggingface.co/v1')
+    model_name = os.getenv('MODEL_NAME', 'meta-llama/Llama-3-8b-instruct')
     
-    # Try HF_TOKEN first, then API_KEY
-    api_key = os.getenv('HF_TOKEN') or os.getenv('API_KEY')
+    # Only use HF_TOKEN (remove API_KEY fallback)
+    hf_token = os.getenv('HF_TOKEN')
     
-    if not api_key:
-        raise ValueError("Neither HF_TOKEN nor API_KEY environment variables are set")
+    if not hf_token:
+        raise ValueError("HF_TOKEN environment variable is required")
     
-    client_kwargs = {'api_key': api_key}
-    if api_base_url:
-        client_kwargs['base_url'] = api_base_url
+    # Initialize OpenAI client with base_url and api_key
+    client = OpenAI(
+        api_key=hf_token,
+        base_url=api_base_url
+    )
     
-    return OpenAI(**client_kwargs), model_name
+    return client, model_name
 
 def create_adaptive_prompt(observation: Dict[str, Any], step_type: str, agent_state: AgentState) -> str:
     """Create adaptive prompt based on agent learning state."""
@@ -159,6 +161,7 @@ def run_llm_inference(client: OpenAI, model_name: str, prompt: str) -> str:
 
 def run_baseline(env: CodeGuardianEnv) -> float:
     """Run a simple non-learning baseline agent for comparison."""
+    print("START: Running baseline agent...")
     logger.info("[BASELINE START] Running baseline agent...")
     
     obs = env.reset()
@@ -185,6 +188,7 @@ def run_baseline(env: CodeGuardianEnv) -> float:
         
         logger.info(f"[BASELINE STEP] {step_count}: {step_type} -> reward={reward.value:+.3f}")
     
+    print(f"END: Baseline complete - Total reward: {total_reward:.3f}, Steps: {step_count}")
     logger.info(f"[BASELINE END] Total reward: {total_reward:.3f}, Steps: {step_count}")
     return total_reward
 
@@ -234,6 +238,7 @@ def run_episode(env: CodeGuardianEnv, client: OpenAI, model_name: str, agent_sta
     episode_history = []
     step_count = 0
     
+    print(f"START: Episode {agent_state.episode + 1} - model={model_name} task_id={obs.task_id} difficulty={obs.difficulty}")
     logger.info(f"[EPISODE START] Episode {agent_state.episode + 1} - model={model_name} task_id={obs.task_id} difficulty={obs.difficulty}")
     
     while not env.done:
@@ -242,6 +247,8 @@ def run_episode(env: CodeGuardianEnv, client: OpenAI, model_name: str, agent_sta
         step_type = step_types[current_step]
         
         step_count += 1
+        print(f"STEP: {step_type}")
+        print(f"STEP: {step_count} - {step_type} - task_id={obs.task_id}")
         logger.info(f"[STEP] step_number={step_count} step_type={step_type} task_id={obs.task_id}")
         
         # Create adaptive prompt
@@ -303,6 +310,7 @@ def run_episode(env: CodeGuardianEnv, client: OpenAI, model_name: str, agent_sta
             break
     
     # Final episode log
+    print(f"END: Episode {agent_state.episode + 1} - final_score={episode_reward:.3f} total_steps={step_count} strategy={agent_state.prompt_strategy}")
     logger.info(f"[EPISODE END] Episode {agent_state.episode + 1} - final_score={episode_reward:.3f} total_steps={step_count} strategy={agent_state.prompt_strategy}")
     
     return episode_reward, episode_history
@@ -315,6 +323,7 @@ def train_agent(num_episodes: int = 10) -> AgentState:
         agent_state = AgentState()
         all_episodes = []
         
+        print(f"START: RL training for {num_episodes} episodes with model: {model_name}")
         logger.info(f"[TRAINING START] Starting RL training for {num_episodes} episodes with model: {model_name}")
         
         # RUN BASELINE FIRST
@@ -357,8 +366,9 @@ def train_agent(num_episodes: int = 10) -> AgentState:
                 break
         
         # PERFORMANCE COMPARISON
+        print("START: Training summary analysis")
         logger.info(f"\n{'='*60}")
-        logger.info("[TRAINING SUMMARY]")
+        logger.info(f"[TRAINING SUMMARY]")
         logger.info(f"{'='*60}")
         
         # Calculate improvement over baseline
@@ -378,12 +388,16 @@ def train_agent(num_episodes: int = 10) -> AgentState:
             
             # Performance classification
             if improvement > 50:
+                print("🚀 EXCELLENT: Significant improvement over baseline!")
                 logger.info("🚀 EXCELLENT: Significant improvement over baseline!")
             elif improvement > 20:
+                print("✅ GOOD: Measurable improvement achieved")
                 logger.info("✅ GOOD: Measurable improvement achieved")
             elif improvement > 0:
+                print("📈 POSITIVE: Some improvement over baseline")
                 logger.info("📈 POSITIVE: Some improvement over baseline")
             else:
+                print("⚠️  NO IMPROVEMENT: RL agent not better than baseline")
                 logger.info("⚠️  NO IMPROVEMENT: RL agent not better than baseline")
         
         return agent_state
@@ -395,11 +409,17 @@ def train_agent(num_episodes: int = 10) -> AgentState:
 def main():
     """Main inference loop with RL training."""
     try:
+        # Required log at beginning
+        print("START: inference")
+        
         # Run RL training
         agent_state = train_agent(num_episodes=10)
+        print("END: done")
+        logger.info("END: RL training completed successfully")
         return 0
         
     except Exception as e:
+        print(f"END: done")
         logger.error(f"Inference failed: {e}")
         return 1