initial-setup

README_inference.md

@@ -0,0 +1,93 @@
+# CodeGuardian Inference Script
+
+This script runs an LLM agent against the CodeGuardian multi-step code review environment.
+
+## Setup
+
+### Environment Variables
+
+Set the following environment variables before running:
+
+```bash
+# Required: API endpoint (for OpenAI-compatible APIs)
+export API_BASE_URL="https://api.openai.com/v1"
+
+# Required: Model name to use
+export MODEL_NAME="gpt-4"
+
+# Required: Authentication token (HF_TOKEN takes precedence over API_KEY)
+export HF_TOKEN="your-huggingface-token"
+# OR
+export API_KEY="your-openai-api-key"
+```
+
+### Dependencies
+
+Install required packages:
+
+```bash
+pip install openai pydantic
+```
+
+## Usage
+
+Run the inference script:
+
+```bash
+python inference.py
+```
+
+## What It Does
+
+The script:
+
+1. **Initializes** the CodeGuardian environment and OpenAI client
+2. **Loops through** all tasks in the environment (3 tasks × 3 steps each = 9 total steps)
+3. **For each step**:
+   - Step 0: Detect issues in the code
+   - Step 1: Assess risk level (low/medium/high)
+   - Step 2: Suggest corrected code
+4. **Sends structured prompts** to the LLM with the current observation
+5. **Parses JSON responses** safely with fallback handling
+6. **Executes actions** in the environment and accumulates rewards
+7. **Logs progress** including rewards, feedback, and running totals
+
+## Output Format
+
+```
+2026-04-06 02:26:51,542 - INFO - Starting inference with model: gpt-4
+2026-04-06 02:26:51,561 - INFO - Initial task: Return sum of two numbers
+2026-04-06 02:26:51,561 - INFO - --- Step 1: DETECT ---
+2026-04-06 02:26:51,562 - INFO - LLM response: {"detected_issues": ["wrong operator"], ...}
+2026-04-06 02:26:51,564 - INFO - Action type: detect
+2026-04-06 02:26:51,564 - INFO - Reward: +0.30
+2026-04-06 02:26:51,564 - INFO - Feedback: []
+2026-04-06 02:26:51,565 - INFO - Running total: +0.30
+...
+2026-04-06 02:26:51,586 - INFO - === Inference Complete ===
+2026-04-06 02:26:51,586 - INFO - Final score: 0.00
+2026-04-06 02:26:51,586 - INFO - Total steps: 9
+```
+
+## Reward Structure
+
+- **Detect step**: +0.3 (correct) / -0.2 (incorrect)
+- **Risk step**: +0.2 (correct) / -0.2 (incorrect)  
+- **Fix step**: +0.5 (correct) / -0.3 (incorrect)
+- **Invalid action order**: -0.5
+
+## Error Handling
+
+The script is robust against:
+- Invalid JSON responses (falls back to empty responses)
+- API errors (falls back to empty responses)
+- Missing fields (provides sensible defaults)
+- Infinite loops (50-step safety limit)
+
+## Expected Performance
+
+- **Perfect score**: 3.0 (all 3 tasks completed perfectly)
+- **Random baseline**: ~0.0 (rewards balance out)
+- **Typical LLM**: 0.5-2.0 (depending on model capability)
+
+The script provides a stable, reproducible baseline for evaluating LLM performance on multi-step code review tasks.

env/environment.py

@@ -0,0 +1,90 @@
+from env.models import Observation, Action
+from env.tasks import TASKS
+from env.grader import grade_action
+
+class CodeGuardianEnv:
+    def __init__(self):
+        self.tasks = TASKS
+        self.current_task_index = 0
+        self.current_step = 0  # 0 = detect, 1 = risk, 2 = fix
+        self.partial_score = 0
+        self.done = False
+
+    def reset(self):
+        self.current_task_index = 0
+        self.current_step = 0
+        self.done = False
+        self.partial_score = 0
+        return self._get_obs()
+
+    def _get_obs(self):
+        task = self.tasks[self.current_task_index]
+        return Observation(
+            code=task["code"],
+            intent=task["intent"],
+            constraints=task["constraints"]
+        )
+
+    def step(self, action: Action):
+        task = self.tasks[self.current_task_index]
+        expected = task["expected"]
+        reward = 0
+        feedback = []
+        
+        # Check if action order is valid
+        if (self.current_step == 0 and action.action_type != "detect") or \
+           (self.current_step == 1 and action.action_type != "risk") or \
+           (self.current_step == 2 and action.action_type != "fix"):
+            reward = -0.5
+            feedback.append(f"Invalid action order: expected {['detect', 'risk', 'fix'][self.current_step]}, got {action.action_type}")
+        else:
+            # Step 0: detect issues
+            if self.current_step == 0:
+                if action.detected_issues and any(expected["issue"] in issue.lower() for issue in action.detected_issues):
+                    reward = 0.3
+                else:
+                    reward = -0.2
+                    feedback.append("Issue not correctly identified")
+                self.current_step += 1
+            
+            # Step 1: assign risk level
+            elif self.current_step == 1:
+                if action.risk_level and action.risk_level.lower() == expected["risk"]:
+                    reward = 0.2
+                else:
+                    reward = -0.2
+                    feedback.append("Incorrect risk level")
+                self.current_step += 1
+            
+            # Step 2: suggest fix
+            elif self.current_step == 2:
+                if action.suggested_fix and expected["fix_contains"] in action.suggested_fix.lower():
+                    reward = 0.5
+                else:
+                    reward = -0.3
+                    feedback.append("Fix is incorrect or missing")
+                
+                # Move to next task and reset current_step
+                self.current_task_index += 1
+                self.current_step = 0
+                if self.current_task_index >= len(self.tasks):
+                    self.done = True
+        
+        self.partial_score += reward
+        
+        next_obs = self._get_obs() if not self.done else None
+        
+        info = {
+            "feedback": feedback,
+            "total_score": self.partial_score
+        }
+        
+        return next_obs, reward, self.done, info
+
+    def state(self):
+        return {
+            "current_task": self.current_task_index,
+            "current_step": self.current_step,
+            "partial_score": self.partial_score,
+            "done": self.done
+        }

env/grader.py

@@ -0,0 +1,23 @@
+def grade_action(action, expected):
+    score = 0
+    feedback = []
+
+    # Issue detection
+    if any(expected["issue"] in issue.lower() for issue in action.detected_issues):
+        score += 0.3
+    else:
+        feedback.append("Issue not correctly identified")
+
+    # Risk level
+    if action.risk_level.lower() == expected["risk"]:
+        score += 0.2
+    else:
+        feedback.append("Incorrect risk level")
+
+    # Fix correctness
+    if expected["fix_contains"] in action.suggested_fix.lower():
+        score += 0.5
+    else:
+        feedback.append("Fix is incorrect or missing")
+
+    return min(score, 1.0), "; ".join(feedback)

env/models.py

@@ -0,0 +1,20 @@
+from pydantic import BaseModel
+from typing import List, Optional
+
+# What agent sees
+class Observation(BaseModel):
+    code: str
+    intent: str
+    constraints: List[str]
+
+# What agent sends
+class Action(BaseModel):
+    action_type: str  # "detect", "risk", "fix"
+    detected_issues: Optional[List[str]] = None
+    risk_level: Optional[str] = None  # "low", "medium", "high"
+    suggested_fix: Optional[str] = None
+
+# What env returns
+class Reward(BaseModel):
+    score: float
+    feedback: str

env/tasks.py

@@ -0,0 +1,42 @@
+TASKS = [
+    {
+        "id": "easy",
+        "code": "def add(a, b): return a - b",
+        "intent": "Return sum of two numbers",
+        "constraints": ["correct logic"],
+        "expected": {
+            "issue": "wrong operator",
+            "risk": "low",
+            "fix_contains": "+"
+        }
+    },
+    {
+        "id": "medium",
+        "code": """def login(user, password):
+    if user == "admin" and password == "1234":
+        return True""",
+        "intent": "Authenticate user securely",
+        "constraints": ["no hardcoding", "security"],
+        "expected": {
+            "issue": "hardcoded credentials",
+            "risk": "high",
+            "fix_contains": "hashed"
+        }
+    },
+    {
+        "id": "hard",
+        "code": """def process(data):
+    result = []
+    for i in range(len(data)):
+        if data[i] % 2 == 0:
+            result.append(data[i])
+    return result""",
+        "intent": "Return even numbers efficiently",
+        "constraints": ["efficient", "clean code"],
+        "expected": {
+            "issue": "inefficient loop",
+            "risk": "low",
+            "fix_contains": "list comprehension"
+        }
+    }
+]

inference.py

@@ -0,0 +1,206 @@
+#!/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
+
+# Setup logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+def setup_openai_client() -> OpenAI:
+    """Initialize OpenAI client from environment variables."""
+    api_base_url = os.getenv('API_BASE_URL')
+    model_name = os.getenv('MODEL_NAME', 'gpt-3.5-turbo')
+    
+    # Try HF_TOKEN first, then API_KEY
+    api_key = os.getenv('HF_TOKEN') or os.getenv('API_KEY')
+    
+    if not api_key:
+        raise ValueError("Neither HF_TOKEN nor API_KEY environment variables are set")
+    
+    client_kwargs = {'api_key': api_key}
+    if api_base_url:
+        client_kwargs['base_url'] = api_base_url
+    
+    return OpenAI(**client_kwargs), model_name
+
+def create_prompt(observation: Dict[str, Any], step_type: str) -> str:
+    """Create structured prompt for the LLM based on observation and current step."""
+    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'
+    }
+    
+    prompt = f"""You are a professional code reviewer.
+
+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]}
+
+Respond ONLY in JSON format:
+{{
+"detected_issues": [...],
+"risk_level": "...",
+"suggested_fix": "..."
+}}"""
+    
+    return prompt
+
+def parse_llm_response(response_text: str) -> Dict[str, Any]:
+    """Parse LLM response safely 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
+        return {
+            'detected_issues': parsed.get('detected_issues', []),
+            'risk_level': parsed.get('risk_level', 'low'),
+            'suggested_fix': parsed.get('suggested_fix', '')
+        }
+    except (json.JSONDecodeError, AttributeError) as e:
+        logger.warning(f"Failed to parse LLM response: {e}")
+        logger.warning(f"Raw response: {response_text[:200]}...")
+        
+        # Return fallback response
+        return {
+            'detected_issues': [],
+            'risk_level': 'low',
+            'suggested_fix': ''
+        }
+
+def create_action_from_response(step_type: str, parsed_response: Dict[str, Any]) -> Action:
+    """Create Action model from parsed LLM response."""
+    return Action(
+        action_type=step_type,
+        detected_issues=parsed_response['detected_issues'] if step_type == 'detect' else None,
+        risk_level=parsed_response['risk_level'] if step_type == 'risk' else None,
+        suggested_fix=parsed_response['suggested_fix'] if step_type == 'fix' else None
+    )
+
+def run_llm_inference(client: OpenAI, model_name: str, prompt: str) -> str:
+    """Run LLM inference with the 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."},
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0,  # Deterministic
+            max_tokens=300  # Reasonable limit
+        )
+        return response.choices[0].message.content
+    except Exception as e:
+        logger.error(f"LLM inference failed: {e}")
+        raise
+
+def main():
+    """Main inference loop."""
+    try:
+        # Setup
+        client, model_name = setup_openai_client()
+        env = CodeGuardianEnv()
+        obs = env.reset()
+        
+        total_reward = 0
+        step_count = 0
+        
+        logger.info(f"Starting inference with model: {model_name}")
+        logger.info(f"Initial task: {obs.intent}")
+        
+        # Main loop
+        while not env.done:
+            current_step = env.state()['current_step']
+            step_types = ['detect', 'risk', 'fix']
+            step_type = step_types[current_step]
+            
+            step_count += 1
+            logger.info(f"\n--- Step {step_count}: {step_type.upper()} ---")
+            
+            # Create prompt
+            prompt = create_prompt(obs, step_type)
+            
+            # 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 action
+                parsed_response = {'detected_issues': [], 'risk_level': 'low', 'suggested_fix': ''}
+            else:
+                # Parse response
+                parsed_response = parse_llm_response(response_text)
+            
+            # Create action
+            action = create_action_from_response(step_type, parsed_response)
+            
+            # Execute step
+            next_obs, reward, done, info = env.step(action)
+            total_reward += reward
+            
+            # Log results
+            logger.info(f"Action type: {action.action_type}")
+            logger.info(f"Reward: {reward:+.2f}")
+            logger.info(f"Feedback: {info['feedback']}")
+            logger.info(f"Running total: {total_reward:+.2f}")
+            logger.info(f"State: {env.state()}")
+            
+            # Update observation for next step
+            obs = next_obs
+            
+            # Safety check
+            if step_count > 50:  # Prevent infinite loops
+                logger.warning("Maximum steps reached, stopping inference")
+                break
+        
+        # Final output
+        logger.info(f"\n=== Inference Complete ===")
+        logger.info(f"Final score: {total_reward:.2f}")
+        logger.info(f"Total steps: {step_count}")
+        logger.info(f"Final state: {env.state()}")
+        
+    except Exception as e:
+        logger.error(f"Inference failed: {e}")
+        return 1
+    
+    return 0
+
+if __name__ == "__main__":
+    exit(main())

test_env.py

@@ -0,0 +1,48 @@
+from env.environment import CodeGuardianEnv
+from env.models import Action
+
+env = CodeGuardianEnv()
+obs = env.reset()
+
+done = False
+
+step_cycle = ["detect", "risk", "fix"]
+step_index = 0
+
+while not done:
+    print("\n--- NEW STEP ---")
+    print("Code:", obs.code)
+
+    current_action = step_cycle[step_index]
+
+    if current_action == "detect":
+        action = Action(
+            action_type="detect",
+            detected_issues=["wrong operator"],
+            risk_level=None,
+            suggested_fix=None
+        )
+
+    elif current_action == "risk":
+        action = Action(
+            action_type="risk",
+            detected_issues=None,
+            risk_level="low",
+            suggested_fix=None
+        )
+
+    elif current_action == "fix":
+        action = Action(
+            action_type="fix",
+            detected_issues=None,
+            risk_level=None,
+            suggested_fix="return a + b"
+        )
+
+    obs, reward, done, info = env.step(action)
+
+    print("Action:", current_action)
+    print("Reward:", reward)
+    print("Info:", info)
+
+    step_index = (step_index + 1) % 3