agent/llm_agent.py

import json
import textwrap
from openai import OpenAI
from typing import Dict, Any

SYSTEM_PROMPT_BASE = textwrap.dedent(
    """
    You are an expert workforce scheduling agent.
    Your job: assign employees to shifts optimally.
    
    RULES (must follow):
    - Employee skills must include the shift's required_skill
    - Employee must be available on the shift's day (availability[day] == 1)
    - Employee cannot exceed max_hours_per_week
    - No two shifts for the same employee on the same (day, period)
    
    {EFFICIENCY_BIAS}
    
    Respond with ONLY a valid JSON action object — no explanation, no markdown.
    Valid action formats:
      {"action_type": "assign", "employee_id": "emp001", "shift_id": "shf042"}
      {"action_type": "noop"}
    """
).strip()

RECOVERY_BIAS = "The previous approach failed. Try a completely different strategy to assign employees."
EFFICIENCY_BIAS = "Prefer actions that move directly toward task completion in fewer steps. Avoid redundant exploration."

class LLMAgent:
    def __init__(self, client: OpenAI, model_name: str):
        self.client = client
        self.model_name = model_name
        self.temperature = 0.0
        self.max_tokens = 150
        
        self.actions = []
        self.rewards = []
        self.errors = []
        
        # Elite state tracking
        self.best_reward = -float('inf')
        self.last_successful_action = None
        self.task_type = None

    def reset(self):
        self.actions = []
        self.rewards = []
        self.errors = []
        self.best_reward = -float('inf')
        self.last_successful_action = None
        self.task_type = None

    def _infer_task_type(self, obs_dict: Dict[str, Any]) -> str:
        # Task type awareness smart edge
        keys_str = str(obs_dict.keys()).lower()
        if "ui" in keys_str or "element" in keys_str or "screen" in keys_str or "viewport" in keys_str:
            return "navigation"
        return "reasoning"

    def _validate_action_text(self, text: str) -> dict:
        if not text or len(text.strip()) == 0:
            raise ValueError("Empty output from generation.")
        if len(text) > 300: # Fast proxy heuristic for detecting raw text / explanations
            raise ValueError("Output too long, likely contains explanation text.")
            
        clean_text = text.replace("```json", "").replace("```", "").strip()
        try:
            parsed = json.loads(clean_text)
        except json.JSONDecodeError:
            raise ValueError("Invalid JSON format.")
            
        if not isinstance(parsed, dict) or "action_type" not in parsed:
            raise ValueError("Missing 'action_type' required key.")
            
        return parsed

    def generate_action(self, obs_dict: Dict[str, Any], last_reward: float = None, last_error: str = None) -> dict:
        # 1. Edge/Task Check
        if len(self.actions) == 0:
            self.task_type = self._infer_task_type(obs_dict)
            
        # 2. Reward-Aware Tracking / Micro Memory Biasing
        if last_reward is not None:
            self.rewards.append(last_reward)
            if self.best_reward != -float('inf') and last_reward > self.best_reward:
                self.last_successful_action = self.actions[-1] if self.actions else {"action_type": "noop"}
            if last_reward > self.best_reward:
                self.best_reward = last_reward
                
        if last_error:
            self.errors.append(last_error)
            
        # Output Compression: Minimal token structuring
        slim = {
            "unassigned": obs_dict.get("unassigned_shifts", [])[:8],
            "emps": [
                {k: e[k] for k in ("id", "skills", "availability", "assigned_hours", "max_hours_per_week")}
                for e in obs_dict.get("employees", [])
            ],
            "shifts": [
                {k: s[k] for k in ("id", "day", "period", "required_skill", "duration_hours")}
                for s in obs_dict.get("shifts", [])
                if s["id"] in obs_dict.get("unassigned_shifts", [])[:8]
            ],
        }
        
        # State Context Construct
        prompt_injections = []
        
        if self.task_type == "navigation":
            prompt_injections.append("Treat this as visual navigation: Use precise, minimal actions targeting interactive elements.")
        else:
            prompt_injections.append("Treat this as logical reasoning: Use structured rational progression.")
            
        # Strategy Bias based on Performance drop
        if last_reward is not None and last_reward < self.best_reward:
            prompt_injections.append("Avoid strategies that previously reduced reward.")
            
        # Failure Check for Context Switch
        if (len(self.rewards) >= 2 and self.rewards[-1] <= 0.00 and self.rewards[-2] <= 0.00) or (self.errors and self.errors[-1] is not None and self.errors[-1] != "null"):
            prompt_injections.append(RECOVERY_BIAS)
            
        base_prompt = SYSTEM_PROMPT_BASE.replace("{EFFICIENCY_BIAS}", EFFICIENCY_BIAS)
        compiled_prompt = base_prompt + "\n\nCRITICAL CONTEXT:\n" + "\n".join(prompt_injections)
        
        user_payload = {
            "CURRENT STATE": slim,
            "PREVIOUS ACTIONS": self.actions[-3:], # Short context to compress
        }
        if self.last_successful_action:
            user_payload["BEST PAST DIRECTIVE"] = self.last_successful_action

        messages = [
            {"role": "system", "content": compiled_prompt},
            {"role": "user", "content": json.dumps(user_payload)},
        ]
        
        action_dict = {"action_type": "noop"} # Safe default mapping
        
        # Primary Action Generation Block
        try:
            completion = self.client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                temperature=self.temperature,
                max_tokens=self.max_tokens,
                stream=False,
            )
            attempt_text = completion.choices[0].message.content or ""
            action_dict = self._validate_action_text(attempt_text)  # Strict validation layer
            
        except Exception as initial_err:
            # Confidence-Based Retry System -> Single fast retry loop with strict penalty prompts
            try:
                retry_messages = messages + [
                    {"role": "assistant", "content": str(attempt_text) if 'attempt_text' in locals() else "Error"},
                    {"role": "user", "content": "Return ONLY a valid executable action JSON constraint. No explanation."}
                ]
                completion_retry = self.client.chat.completions.create(
                    model=self.model_name,
                    messages=retry_messages,
                    temperature=0.0,
                    max_tokens=self.max_tokens,
                    stream=False,
                )
                action_dict = self._validate_action_text(completion_retry.choices[0].message.content or "")
            except Exception:
                # Safe Fallback Action - Double failure defaults seamlessly safely without loop crash
                action_dict = {"action_type": "noop"}
                
        # Micro-Memory tracking caps
        self.actions.append(action_dict)
        if len(self.actions) > 5: self.actions.pop(0)
        if len(self.rewards) > 5: self.rewards.pop(0)
        if len(self.errors) > 5: self.errors.pop(0)
            
        return action_dict