Spaces:
Sleeping
Sleeping
| 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 | |