minor fixes

inference.py

@@ -20,22 +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', 'https://router.huggingface.co/v1')
-    model_name = os.getenv('MODEL_NAME', 'meta-llama/Llama-3-8b-instruct')
     
-    # Only use HF_TOKEN (remove API_KEY fallback)
-    hf_token = os.getenv('HF_TOKEN')
-    
-    if not hf_token:
+    print("STEP: Initializing OpenAI client")
+
+    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")
+    HF_TOKEN = os.getenv("HF_TOKEN")
+
+    if not HF_TOKEN:
         raise ValueError("HF_TOKEN environment variable is required")
-    
-    # Initialize OpenAI client with base_url and api_key
+
     client = OpenAI(
-        api_key=hf_token,
-        base_url=api_base_url
+        api_key=HF_TOKEN,
+        base_url=API_BASE_URL
     )
-    
-    return client, 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."""
@@ -142,11 +142,11 @@ def create_action_from_response(step_type: str, response: AgentResponse) -> Acti
         suggested_fix=response.fix if step_type == 'fix' else None
     )
 
-def run_llm_inference(client: OpenAI, model_name: str, prompt: str) -> str:
+def run_llm_inference(client: OpenAI, MODEL_NAME: str, prompt: str) -> str:
     """Run LLM inference with given prompt."""
     try:
         response = client.chat.completions.create(
-            model=model_name,
+            model=MODEL_NAME,
             messages=[
                 {"role": "system", "content": "You are a professional code reviewer. Always respond in valid JSON format with confidence and reasoning."},
                 {"role": "user", "content": prompt}
@@ -231,15 +231,15 @@ def update_agent_state(agent_state: AgentState, episode_reward: float, feedback:
     
     return agent_state
 
-def run_episode(env: CodeGuardianEnv, client: OpenAI, model_name: str, agent_state: AgentState) -> tuple[float, List[Dict]]:
+def run_episode(env: CodeGuardianEnv, client: OpenAI, MODEL_NAME: str, agent_state: AgentState) -> tuple[float, List[Dict]]:
     """Run a single episode and return total reward with episode history."""
     obs = env.reset()
     episode_reward = 0.0
     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}")
+    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:
         current_step = env.state()['current_step']
@@ -255,7 +255,7 @@ def run_episode(env: CodeGuardianEnv, client: OpenAI, model_name: str, agent_sta
         
         # Run LLM inference
         try:
-            response_text = run_llm_inference(client, model_name, prompt)
+            response_text = run_llm_inference(client, MODEL_NAME, prompt)
             logger.info(f"LLM response: {response_text[:100]}...")
         except Exception as e:
             logger.error(f"LLM inference failed: {e}")
@@ -318,12 +318,12 @@ def train_agent(num_episodes: int = 10) -> AgentState:
     """Train agent over multiple episodes with policy improvement."""
     try:
         # Setup
-        client, model_name = setup_openai_client()
+        client, MODEL_NAME = setup_openai_client()
         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}")
+        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
         logger.info("\n{'='*60}")
@@ -344,7 +344,7 @@ def train_agent(num_episodes: int = 10) -> AgentState:
             logger.info(f"{'='*50}")
             
             env = CodeGuardianEnv()
-            episode_reward, episode_history = run_episode(env, client, model_name, agent_state)
+            episode_reward, episode_history = run_episode(env, client, MODEL_NAME, agent_state)
             all_episodes.append({
                 "episode": episode + 1,
                 "reward": episode_reward,