Spaces:

testtej44
/

Code_man

Runtime error

Code_man / inference.py

Tejasva

fix:minor

3341e6c 7 months ago

21.2 kB

#!/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__)

# Module-level environment variables
LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")

DEFAULT_API_BASE_URL = "https://router.huggingface.co/v1"
DEFAULT_MODEL_NAME = "meta-llama/Meta-Llama-3-8B-Instruct"
HF_FALLBACK_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"


def _hf_auth_tokens() -> Optional[str]:
    return os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")


def _openai_auth_tokens() -> Optional[str]:
    return os.getenv("OPENAI_API_KEY") or os.getenv("API_KEY")


def _looks_like_huggingface_endpoint(url: str) -> bool:
    u = (url or "").lower()
    return any(x in u for x in ("huggingface", "hf.co", "router.huggingface"))


def _looks_like_openai_endpoint(url: str) -> bool:
    return "openai.com" in (url or "").lower()


def _looks_like_hf_hub_model(model: str) -> bool:
    """Hugging Face router / Inference typically uses org/model ids."""
    return "/" in (model or "").strip()


def _looks_like_openai_model(model: str) -> bool:
    m = (model or "").strip().lower()
    if not m or "/" in m:
        return False
    prefixes = (
        "gpt-",
        "o1",
        "o3",
        "o4",
        "ft:",
        "davinci",
        "text-embedding",
        "chatgpt-",
    )
    return any(m.startswith(p) for p in prefixes)


def _resolve_api_key(api_base_url: str, model_name: str) -> str:
    """
    Pick API key for an OpenAI-compatible client: Hugging Face token and/or OpenAI key,
    ordered from URL + model hints, with 'or' fallbacks so either stack can work.
    """
    hf_key = _hf_auth_tokens()
    oai_key = _openai_auth_tokens()

    url_hf = _looks_like_huggingface_endpoint(api_base_url)
    url_oai = _looks_like_openai_endpoint(api_base_url)
    model_hf = _looks_like_hf_hub_model(model_name)
    model_oai = _looks_like_openai_model(model_name)

    prefer_hf = url_hf or model_hf
    prefer_oai = url_oai or model_oai

    key: Optional[str] = None
    if prefer_hf and not prefer_oai:
        key = hf_key or oai_key
    elif prefer_oai and not prefer_hf:
        key = oai_key or hf_key
    elif prefer_hf and prefer_oai:
        if url_hf and not url_oai:
            key = hf_key or oai_key
        elif url_oai and not url_hf:
            key = oai_key or hf_key
        else:
            key = hf_key or oai_key
    else:
        key = hf_key or oai_key

    if not key:
        raise ValueError(
            "Set authentication for your provider: "
            "(HF_TOKEN or HUGGINGFACEHUB_API_TOKEN) for Hugging Face, "
            "or (OPENAI_API_KEY or API_KEY) for OpenAI"
        )
    return key


def _fallback_llm_json_string(reason: str = "LLM unavailable") -> str:
    return json.dumps(
        {
            "issues_detected": [],
            "risk_level": "low",
            "fix": "",
            "confidence": 0.1,
            "reasoning": reason,
            "reward_breakdown": {"detect": 0.0, "risk": 0.0, "fix": 0.0},
        }
    )


def setup_openai_client() -> tuple[OpenAI, str]:
    """Initialize OpenAI client from environment variables."""

    print("STEP: Initializing OpenAI client")

    api_base_url = os.getenv("API_BASE_URL", DEFAULT_API_BASE_URL)
    model_name = os.getenv("MODEL_NAME", DEFAULT_MODEL_NAME)
    api_key = _resolve_api_key(api_base_url, model_name)

    hf_key = _hf_auth_tokens()
    oai_key = _openai_auth_tokens()
    if api_key == hf_key and hf_key:
        auth_label = "Hugging Face token (HF_TOKEN / HUGGINGFACEHUB_API_TOKEN)"
    elif api_key == oai_key and oai_key:
        auth_label = "OpenAI key (OPENAI_API_KEY / API_KEY)"
    else:
        auth_label = "resolved API key"

    print(f"DEBUG: API_BASE_URL={api_base_url}")
    print(f"DEBUG: MODEL_NAME={model_name}")
    print(f"DEBUG: auth={auth_label}")
    logger.info(
        "OpenAI-compatible client: API_BASE_URL=%s MODEL_NAME=%s auth=%s",
        api_base_url,
        model_name,
        auth_label,
    )

    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 _is_model_missing_error(exc: BaseException) -> bool:
    msg = str(exc).lower()
    if "model_not_found" in msg or "model not found" in msg:
        return True
    body = getattr(exc, "body", None)
    if isinstance(body, str) and "model_not_found" in body.lower():
        return True
    if isinstance(body, dict):
        err = body.get("error")
        if isinstance(err, dict) and err.get("code") == "model_not_found":
            return True
    return False


def run_llm_inference(client: OpenAI, model_name: str, prompt: str) -> str:
    """Run LLM inference; never raises — returns JSON text or a safe fallback."""
    models_to_try: List[str] = [model_name]
    if HF_FALLBACK_MODEL and HF_FALLBACK_MODEL not in models_to_try:
        models_to_try.append(HF_FALLBACK_MODEL)

    last_error: Optional[BaseException] = None
    for m in models_to_try:
        try:
            response = client.chat.completions.create(
                model=m,
                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,
                max_tokens=500,
            )
            content = response.choices[0].message.content
            if content is None or not str(content).strip():
                logger.error("LLM returned empty content for model=%s", m)
                last_error = ValueError("empty LLM content")
                continue
            if m != model_name:
                logger.info("LLM succeeded using fallback model=%s (primary=%s failed)", m, model_name)
            return str(content)
        except Exception as e:
            last_error = e
            logger.error(
                "LLM inference failed for model=%s: %s",
                m,
                e,
                exc_info=not _is_model_missing_error(e),
            )
            if _is_model_missing_error(e):
                logger.error("model_not_found for model=%s; will try next model if available", m)
            continue

    logger.error("All LLM models exhausted; last_error=%s", last_error)
    return _fallback_llm_json_string("LLM call failed — using fallback response")

def run_baseline(env: CodeGuardianEnv) -> float:
    """Run a simple non-learning baseline agent for comparison."""
    print("STEP: 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"STEP: Baseline complete - 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"STEP: Episode {agent_state.episode + 1} started - 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_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 (never raises)
        response_text = run_llm_inference(client, model_name, prompt)
        logger.info("LLM response (truncated): %s...", response_text[:100])
        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"STEP: Episode {agent_state.episode + 1} complete - 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"STEP: RL training started - episodes={num_episodes} 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("STEP: 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("STEP: EXCELLENT - significant improvement over baseline")
            elif improvement > 20:
                print("STEP: GOOD - measurable improvement achieved")
            elif improvement > 0:
                print("STEP: POSITIVE - some improvement over baseline")
            else:
                print("STEP: 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:
        print("START")
        
        LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
        if LOCAL_IMAGE_NAME:
            print(f"STEP: inference using Docker image {LOCAL_IMAGE_NAME}")
        else:
            print("STEP: Starting inference")
        
        agent_state = train_agent(num_episodes=10)
        return 0
        
    except Exception as e:
        logger.error(f"Inference failed: {e}")
        return 1
    finally:
        print("END")

if __name__ == "__main__":
    exit(main())

	#!/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__)

	# Module-level environment variables
	LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")

	DEFAULT_API_BASE_URL = "https://router.huggingface.co/v1"
	DEFAULT_MODEL_NAME = "meta-llama/Meta-Llama-3-8B-Instruct"
	HF_FALLBACK_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"


	def _hf_auth_tokens() -> Optional[str]:
	return os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")


	def _openai_auth_tokens() -> Optional[str]:
	return os.getenv("OPENAI_API_KEY") or os.getenv("API_KEY")


	def _looks_like_huggingface_endpoint(url: str) -> bool:
	u = (url or "").lower()
	return any(x in u for x in ("huggingface", "hf.co", "router.huggingface"))


	def _looks_like_openai_endpoint(url: str) -> bool:
	return "openai.com" in (url or "").lower()


	def _looks_like_hf_hub_model(model: str) -> bool:
	"""Hugging Face router / Inference typically uses org/model ids."""
	return "/" in (model or "").strip()


	def _looks_like_openai_model(model: str) -> bool:
	m = (model or "").strip().lower()
	if not m or "/" in m:
	return False
	prefixes = (
	"gpt-",
	"o1",
	"o3",
	"o4",
	"ft:",
	"davinci",
	"text-embedding",
	"chatgpt-",
	)
	return any(m.startswith(p) for p in prefixes)


	def _resolve_api_key(api_base_url: str, model_name: str) -> str:
	"""
	Pick API key for an OpenAI-compatible client: Hugging Face token and/or OpenAI key,
	ordered from URL + model hints, with 'or' fallbacks so either stack can work.
	"""
	hf_key = _hf_auth_tokens()
	oai_key = _openai_auth_tokens()

	url_hf = _looks_like_huggingface_endpoint(api_base_url)
	url_oai = _looks_like_openai_endpoint(api_base_url)
	model_hf = _looks_like_hf_hub_model(model_name)
	model_oai = _looks_like_openai_model(model_name)

	prefer_hf = url_hf or model_hf
	prefer_oai = url_oai or model_oai

	key: Optional[str] = None
	if prefer_hf and not prefer_oai:
	key = hf_key or oai_key
	elif prefer_oai and not prefer_hf:
	key = oai_key or hf_key
	elif prefer_hf and prefer_oai:
	if url_hf and not url_oai:
	key = hf_key or oai_key
	elif url_oai and not url_hf:
	key = oai_key or hf_key
	else:
	key = hf_key or oai_key
	else:
	key = hf_key or oai_key

	if not key:
	raise ValueError(
	"Set authentication for your provider: "
	"(HF_TOKEN or HUGGINGFACEHUB_API_TOKEN) for Hugging Face, "
	"or (OPENAI_API_KEY or API_KEY) for OpenAI"
	)
	return key


	def _fallback_llm_json_string(reason: str = "LLM unavailable") -> str:
	return json.dumps(
	{
	"issues_detected": [],
	"risk_level": "low",
	"fix": "",
	"confidence": 0.1,
	"reasoning": reason,
	"reward_breakdown": {"detect": 0.0, "risk": 0.0, "fix": 0.0},
	}
	)


	def setup_openai_client() -> tuple[OpenAI, str]:
	"""Initialize OpenAI client from environment variables."""

	print("STEP: Initializing OpenAI client")

	api_base_url = os.getenv("API_BASE_URL", DEFAULT_API_BASE_URL)
	model_name = os.getenv("MODEL_NAME", DEFAULT_MODEL_NAME)
	api_key = _resolve_api_key(api_base_url, model_name)

	hf_key = _hf_auth_tokens()
	oai_key = _openai_auth_tokens()
	if api_key == hf_key and hf_key:
	auth_label = "Hugging Face token (HF_TOKEN / HUGGINGFACEHUB_API_TOKEN)"
	elif api_key == oai_key and oai_key:
	auth_label = "OpenAI key (OPENAI_API_KEY / API_KEY)"
	else:
	auth_label = "resolved API key"

	print(f"DEBUG: API_BASE_URL={api_base_url}")
	print(f"DEBUG: MODEL_NAME={model_name}")
	print(f"DEBUG: auth={auth_label}")
	logger.info(
	"OpenAI-compatible client: API_BASE_URL=%s MODEL_NAME=%s auth=%s",
	api_base_url,
	model_name,
	auth_label,
	)

	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 _is_model_missing_error(exc: BaseException) -> bool:
	msg = str(exc).lower()
	if "model_not_found" in msg or "model not found" in msg:
	return True
	body = getattr(exc, "body", None)
	if isinstance(body, str) and "model_not_found" in body.lower():
	return True
	if isinstance(body, dict):
	err = body.get("error")
	if isinstance(err, dict) and err.get("code") == "model_not_found":
	return True
	return False


	def run_llm_inference(client: OpenAI, model_name: str, prompt: str) -> str:
	"""Run LLM inference; never raises — returns JSON text or a safe fallback."""
	models_to_try: List[str] = [model_name]
	if HF_FALLBACK_MODEL and HF_FALLBACK_MODEL not in models_to_try:
	models_to_try.append(HF_FALLBACK_MODEL)

	last_error: Optional[BaseException] = None
	for m in models_to_try:
	try:
	response = client.chat.completions.create(
	model=m,
	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,
	max_tokens=500,
	)
	content = response.choices[0].message.content
	if content is None or not str(content).strip():
	logger.error("LLM returned empty content for model=%s", m)
	last_error = ValueError("empty LLM content")
	continue
	if m != model_name:
	logger.info("LLM succeeded using fallback model=%s (primary=%s failed)", m, model_name)
	return str(content)
	except Exception as e:
	last_error = e
	logger.error(
	"LLM inference failed for model=%s: %s",
	m,
	e,
	exc_info=not _is_model_missing_error(e),
	)
	if _is_model_missing_error(e):
	logger.error("model_not_found for model=%s; will try next model if available", m)
	continue

	logger.error("All LLM models exhausted; last_error=%s", last_error)
	return _fallback_llm_json_string("LLM call failed — using fallback response")

	def run_baseline(env: CodeGuardianEnv) -> float:
	"""Run a simple non-learning baseline agent for comparison."""
	print("STEP: 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"STEP: Baseline complete - 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"STEP: Episode {agent_state.episode + 1} started - 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_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 (never raises)
	response_text = run_llm_inference(client, model_name, prompt)
	logger.info("LLM response (truncated): %s...", response_text[:100])
	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"STEP: Episode {agent_state.episode + 1} complete - 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"STEP: RL training started - episodes={num_episodes} 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("STEP: 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("STEP: EXCELLENT - significant improvement over baseline")
	elif improvement > 20:
	print("STEP: GOOD - measurable improvement achieved")
	elif improvement > 0:
	print("STEP: POSITIVE - some improvement over baseline")
	else:
	print("STEP: 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:
	print("START")

	LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
	if LOCAL_IMAGE_NAME:
	print(f"STEP: inference using Docker image {LOCAL_IMAGE_NAME}")
	else:
	print("STEP: Starting inference")

	agent_state = train_agent(num_episodes=10)
	return 0

	except Exception as e:
	logger.error(f"Inference failed: {e}")
	return 1
	finally:
	print("END")

	if __name__ == "__main__":
	exit(main())