minor fixes

inference_backup.py

@@ -0,0 +1,424 @@
+#!/usr/bin/env python3
+
+import os
+import json
+import logging
+from typing import Dict, Any, Optional, List
+
+try:
+    from openai import OpenAI
+except ImportError:
+    print("Error: openai package not found. Install with: pip install openai")
+    exit(1)
+
+from env.environment import CodeGuardianEnv
+from env.models import Action, AgentResponse, AgentState
+
+# Setup logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+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://api.openai.com/v1')
+    model_name = os.getenv('MODEL_NAME', 'gpt-4')
+    
+    # HF_TOKEN without default (required for Hugging Face)
+    hf_token = os.getenv('HF_TOKEN')
+    api_key = hf_token or os.getenv('API_KEY')
+    
+    if not api_key:
+        raise ValueError("Neither HF_TOKEN nor API_KEY environment variables are set")
+    
+    # Initialize OpenAI client with base_url and api_key
+    client = OpenAI(
+        api_key=api_key,
+        base_url=api_base_url
+    )
+    
+    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."""
+    step_names = {
+        'detect': 'Issue Detection',
+        'risk': 'Risk Assessment', 
+        'fix': 'Fix Suggestion'
+    }
+    
+    step_instructions = {
+        'detect': 'list all issues you find in the code',
+        'risk': 'output the risk level (low/medium/high)',
+        'fix': 'provide corrected code that fixes the identified issues'
+    }
+    
+    # Adaptive prompt based on learning
+    learning_context = ""
+    if agent_state.past_failures:
+        learning_context = f"\nLEARNING FROM PAST MISTAKES:\n{chr(10).join(f'- Avoid: {failure}' for failure in agent_state.past_failures[-3:])}"
+    
+    if agent_state.prompt_strategy == "detailed":
+        instruction_detail = " Be thorough and provide detailed reasoning."
+    elif agent_state.prompt_strategy == "focused":
+        instruction_detail = " Focus on most critical issues first."
+    else:
+        instruction_detail = ""
+    
+    prompt = f"""You are a professional code reviewer learning from feedback.
+
+Code:
+{observation.code}
+
+Intent:
+{observation.intent}
+
+Constraints:
+{chr(10).join(f"- {constraint}" for constraint in observation.constraints)}
+
+Step: {step_names[step_type]}
+
+Instructions:
+* {step_instructions[step_type]}{instruction_detail}
+* Provide confidence score (0.0-1.0)
+* Include your reasoning{learning_context}
+
+Respond ONLY in JSON format:
+{{
+  "issues_detected": [...],
+  "risk_level": "...",
+  "fix": "...",
+  "confidence": 0.0,
+  "reasoning": "...",
+  "reward_breakdown": {{
+    "detect": 0.0,
+    "risk": 0.0,
+    "fix": 0.0
+  }}
+}}"""
+    
+    return prompt
+
+def parse_llm_response(response_text: str) -> AgentResponse:
+    """Parse LLM response into structured AgentResponse with fallback."""
+    try:
+        # Clean up response text
+        response_text = response_text.strip()
+        if response_text.startswith('```json'):
+            response_text = response_text[7:]
+        if response_text.endswith('```'):
+            response_text = response_text[:-3]
+        response_text = response_text.strip()
+        
+        parsed = json.loads(response_text)
+        
+        # Ensure all required fields exist with defaults
+        return AgentResponse(
+            issues_detected=parsed.get('issues_detected', []),
+            risk_level=parsed.get('risk_level', 'low'),
+            fix=parsed.get('fix', ''),
+            confidence=max(0.0, min(1.0, parsed.get('confidence', 0.5))),
+            reasoning=parsed.get('reasoning', ''),
+            reward_breakdown=parsed.get('reward_breakdown', {})
+        )
+    except (json.JSONDecodeError, AttributeError, ValueError) as e:
+        logger.warning(f"Failed to parse LLM response: {e}")
+        logger.warning(f"Raw response: {response_text[:200]}...")
+        
+        # Return fallback response
+        return AgentResponse(
+            issues_detected=[],
+            risk_level='low',
+            fix='',
+            confidence=0.1,
+            reasoning='Failed to parse response',
+            reward_breakdown={}
+        )
+
+def create_action_from_response(step_type: str, response: AgentResponse) -> Action:
+    """Create Action model from structured AgentResponse."""
+    return Action(
+        action_type=step_type,
+        detected_issues=response.issues_detected if step_type == 'detect' else None,
+        risk_level=response.risk_level if step_type == 'risk' else None,
+        suggested_fix=response.fix if step_type == 'fix' else None
+    )
+
+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,
+            messages=[
+                {"role": "system", "content": "You are a professional code reviewer. Always respond in valid JSON format with confidence and reasoning."},
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0,  # Deterministic
+            max_tokens=500  # Increased for structured output
+        )
+        return response.choices[0].message.content
+    except Exception as e:
+        logger.error(f"LLM inference failed: {e}")
+        raise
+
+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()
+    total_reward = 0.0
+    step_count = 0
+    
+    while not env.done and step_count < 50:
+        current_step = env.state()['current_step']
+        step_types = ['detect', 'risk', 'fix']
+        step_type = step_types[current_step]
+        
+        # Baseline always outputs minimal responses
+        if step_type == 'detect':
+            action = Action(action_type='detect', detected_issues=[])
+        elif step_type == 'risk':
+            action = Action(action_type='risk', risk_level='low')
+        else:
+            action = Action(action_type='fix', suggested_fix='')
+        
+        next_obs, reward, done = env.step(action)
+        total_reward += reward.value
+        step_count += 1
+        obs = next_obs
+        
+        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
+
+def update_agent_state(agent_state: AgentState, episode_reward: float, feedback: str) -> AgentState:
+    """Update agent state based on episode performance for policy improvement."""
+    agent_state.episode += 1
+    agent_state.total_reward = episode_reward
+    agent_state.improvement_history.append(episode_reward)
+    
+    # Update best score
+    if episode_reward > agent_state.best_score:
+        agent_state.best_score = episode_reward
+        logger.info(f"New best score achieved: {episode_reward:.3f}")
+    
+    # STRONGER POLICY UPDATE LOGIC
+    # Treat poor performance as failure
+    if episode_reward < agent_state.best_score * 0.8:
+        agent_state.past_failures.append(f"Episode {agent_state.episode}: reward={episode_reward:.3f}")
+        # Keep only recent failures
+        agent_state.past_failures = agent_state.past_failures[-5:]
+        logger.info(f"Failure recorded - Episode reward: {episode_reward:.3f} vs Best: {agent_state.best_score:.3f}")
+    
+    # IMPROVED STRATEGY SWITCHING
+    if episode_reward < agent_state.best_score * 0.7:  # Poor performance threshold
+        old_strategy = agent_state.prompt_strategy
+        
+        if agent_state.prompt_strategy == "default":
+            agent_state.prompt_strategy = "detailed"
+            logger.info("Switching strategy: default → detailed")
+        elif agent_state.prompt_strategy == "detailed":
+            agent_state.prompt_strategy = "focused"
+            logger.info("Switching strategy: detailed → focused")
+        else:
+            # Reset to default if focused still fails
+            agent_state.prompt_strategy = "default"
+            logger.info("Switching strategy: focused → default (reset)")
+        
+        if old_strategy != agent_state.prompt_strategy:
+            logger.info(f"Strategy change triggered by low performance: {episode_reward:.3f} < {agent_state.best_score * 0.7:.3f}")
+    
+    return agent_state
+
+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}")
+    
+    while not env.done:
+        current_step = env.state()['current_step']
+        step_types = ['detect', 'risk', 'fix']
+        step_type = step_types[current_step]
+        
+        step_count += 1
+        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
+        prompt = create_adaptive_prompt(obs, step_type, agent_state)
+        
+        # Run LLM inference
+        try:
+            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}")
+            # Use fallback response
+            agent_response = AgentResponse(
+                issues_detected=[],
+                risk_level='low',
+                fix='',
+                confidence=0.1,
+                reasoning='Inference failed',
+                reward_breakdown={}
+            )
+        else:
+            # Parse structured response
+            agent_response = parse_llm_response(response_text)
+        
+        # Create action from response
+        action = create_action_from_response(step_type, agent_response)
+        
+        # Execute step
+        next_obs, reward, done = env.step(action)
+        
+        # CONFIDENCE-BASED REWARD ADJUSTMENT
+        confidence_penalty = (1 - agent_response.confidence) * 0.1
+        adjusted_reward = reward.value - confidence_penalty
+        episode_reward += adjusted_reward
+        
+        logger.info(f"confidence_penalty={confidence_penalty:.3f} adjusted_reward={adjusted_reward:+.3f} confidence={agent_response.confidence:.2f}")
+        
+        # Log step results
+        step_info = {
+            "step": step_count,
+            "type": step_type,
+            "action": action.dict(),
+            "response": agent_response.dict(),
+            "reward": reward.dict(),
+            "adjusted_reward": adjusted_reward,
+            "confidence_penalty": confidence_penalty,
+            "total_reward": episode_reward
+        }
+        episode_history.append(step_info)
+        
+        logger.info(f"action_type={action.action_type} reward={reward.value:+.3f} adjusted_reward={adjusted_reward:+.3f} feedback={reward.feedback} total_reward={episode_reward:+.3f} current_step={env.state()['current_step']} current_task={env.state()['current_task']}")
+        
+        # Update observation for next step
+        obs = next_obs
+        
+        # Safety check
+        if step_count > 50:  # Prevent infinite loops
+            logger.warning("Maximum steps reached, stopping episode")
+            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
+
+def train_agent(num_episodes: int = 10) -> AgentState:
+    """Train agent over multiple episodes with policy improvement."""
+    try:
+        # Setup
+        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}")
+        
+        # RUN BASELINE FIRST
+        logger.info("\n{'='*60}")
+        logger.info("BASELINE EVALUATION")
+        logger.info(f"{'='*60}")
+        baseline_env = CodeGuardianEnv()
+        baseline_score = run_baseline(baseline_env)
+        logger.info(f"Baseline score: {baseline_score:.3f}")
+        
+        # TRAINING LOOP
+        logger.info(f"\n{'='*60}")
+        logger.info("REINFORCEMENT LEARNING")
+        logger.info(f"{'='*60}")
+        
+        for episode in range(num_episodes):
+            logger.info(f"\n{'='*50}")
+            logger.info(f"EPISODE {episode + 1}/{num_episodes}")
+            logger.info(f"{'='*50}")
+            
+            env = CodeGuardianEnv()
+            episode_reward, episode_history = run_episode(env, client, model_name, agent_state)
+            all_episodes.append({
+                "episode": episode + 1,
+                "reward": episode_reward,
+                "history": episode_history
+            })
+            
+            # Update agent state for policy improvement
+            agent_state = update_agent_state(agent_state, episode_reward, f"Episode {episode + 1} completed")
+            
+            logger.info(f"Episode {episode + 1} completed with reward: {episode_reward:.3f}")
+            logger.info(f"Best score so far: {agent_state.best_score:.3f}")
+            logger.info(f"Current strategy: {agent_state.prompt_strategy}")
+            
+            # EARLY STOPPING
+            if episode > 3 and agent_state.best_score > 2.5:
+                logger.info("Early stopping triggered: high performance achieved")
+                logger.info(f"Stopping at episode {episode + 1} with best score {agent_state.best_score:.3f}")
+                break
+        
+        # PERFORMANCE COMPARISON
+        print("START: Training summary analysis")
+        logger.info(f"\n{'='*60}")
+        logger.info(f"[TRAINING SUMMARY]")
+        logger.info(f"{'='*60}")
+        
+        # Calculate improvement over baseline
+        improvement = ((agent_state.best_score - baseline_score) / (abs(baseline_score) + 1e-6)) * 100
+        
+        # Log key metrics
+        logger.info(f"RL Agent Best Score: {agent_state.best_score:.3f}")
+        logger.info(f"Baseline Score: {baseline_score:.3f}")
+        logger.info(f"Performance Improvement: {improvement:.2f}%")
+        
+        # Additional metrics
+        if all_episodes:
+            avg_reward = sum(ep["reward"] for ep in all_episodes) / len(all_episodes)
+            logger.info(f"Average Episode Reward: {avg_reward:.3f}")
+            logger.info(f"Episodes Completed: {len(all_episodes)}")
+            logger.info(f"Final Strategy: {agent_state.prompt_strategy}")
+            
+            # 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
+        
+    except Exception as e:
+        logger.error(f"Training failed: {e}")
+        raise
+
+def main():
+    """Main inference loop with RL training."""
+    try:
+        # Run RL training
+        agent_state = train_agent(num_episodes=10)
+        print("END: RL training completed successfully")
+        logger.info("END: RL training completed successfully")
+        return 0
+        
+    except Exception as e:
+        print(f"END: Inference failed with error: {e}")
+        logger.error(f"Inference failed: {e}")
+        return 1
+
+if __name__ == "__main__":
+    exit(main())

requirements.txt

@@ -5,3 +5,4 @@ numpy>=1.24.0
 python-dotenv>=1.0.0
 flask>=3.0.0
 gradio>=4.0.0
+plotly