"""Sample LLM agent for the HR Productivity Environment using Claude API.""" from __future__ import annotations import json import logging from typing import Any, Dict, List, Optional import anthropic from agent.prompts import PHASE_PROMPTS, QUARTER_SUMMARY_PROMPT, SYSTEM_PROMPT logger = logging.getLogger(__name__) class HRAgent: """Claude-based HR agent with phase-aware prompting and memory management.""" def __init__( self, model: str = "claude-sonnet-4-6", max_tokens: int = 2048, ): self.client = anthropic.Anthropic() self.model = model self.max_tokens = max_tokens self.quarter_summaries: List[str] = [] self.current_quarter: int = 1 self.conversation_history: List[Dict[str, Any]] = [] self.max_history_messages = 20 # Rolling window def decide(self, observation: Dict[str, Any]) -> Dict[str, Any]: """Given an observation, decide the next action. Args: observation: The observation dict from the environment. Returns: Action dict to send to the environment. """ phase = observation.get("current_phase", "scanning") quarter = observation.get("current_quarter", 1) # Detect quarter change if quarter != self.current_quarter: self._summarize_quarter() self.current_quarter = quarter self.conversation_history = [] # Reset conversation for new quarter # Build the prompt user_message = self._build_user_message(observation, phase) self.conversation_history.append({"role": "user", "content": user_message}) # Trim history if too long if len(self.conversation_history) > self.max_history_messages: self.conversation_history = self.conversation_history[-self.max_history_messages:] # Build system prompt with memory system = self._build_system_prompt(phase) try: response = self.client.messages.create( model=self.model, max_tokens=self.max_tokens, system=system, messages=self.conversation_history, ) assistant_text = response.content[0].text self.conversation_history.append({"role": "assistant", "content": assistant_text}) # Parse action from response action = self._parse_action(assistant_text) return action except Exception as e: logger.error(f"Agent error: {e}") # Fallback: advance phase or query return self._fallback_action(phase, observation) def _build_system_prompt(self, phase: str) -> str: """Build system prompt with phase guidance and quarter memory.""" parts = [SYSTEM_PROMPT] if self.quarter_summaries: parts.append("\n## Previous Quarter Summaries:") for i, summary in enumerate(self.quarter_summaries, 1): parts.append(f"\n### Q{i}:\n{summary}") parts.append(f"\n{PHASE_PROMPTS.get(phase, '')}") return "\n".join(parts) def _build_user_message(self, observation: Dict[str, Any], phase: str) -> str: """Build user message from observation.""" parts = [f"**Quarter {observation.get('current_quarter', '?')} — {phase.upper()} phase**"] parts.append(f"Step {observation.get('step_in_quarter', 0)} in quarter") parts.append(f"\n**Message:** {observation.get('message', 'No message')}") if observation.get("warnings"): parts.append(f"\n**Warnings:** {', '.join(observation['warnings'])}") if observation.get("data"): # Truncate large data data_str = json.dumps(observation["data"], indent=2, default=str) if len(data_str) > 3000: data_str = data_str[:3000] + "\n... (truncated)" parts.append(f"\n**Data:**\n```json\n{data_str}\n```") available = observation.get("available_actions", []) parts.append(f"\n**Available actions:** {', '.join(available)}") parts.append("\nRespond with a JSON action object.") return "\n".join(parts) def _parse_action(self, text: str) -> Dict[str, Any]: """Extract JSON action from the LLM response.""" # Try to find JSON block if "```json" in text: start = text.index("```json") + 7 end = text.index("```", start) json_str = text[start:end].strip() elif "```" in text: start = text.index("```") + 3 end = text.index("```", start) json_str = text[start:end].strip() elif "{" in text: start = text.index("{") # Find matching closing brace depth = 0 for i, c in enumerate(text[start:], start): if c == "{": depth += 1 elif c == "}": depth -= 1 if depth == 0: json_str = text[start:i + 1] break else: json_str = text[start:] else: return {"action_type": "advance_phase", "rationale": "Could not parse action, advancing"} try: action = json.loads(json_str) # Ensure action_type exists if "action_type" not in action: action["action_type"] = "advance_phase" return action except json.JSONDecodeError: return {"action_type": "advance_phase", "rationale": "JSON parse error, advancing"} def _fallback_action(self, phase: str, observation: Dict[str, Any]) -> Dict[str, Any]: """Generate a safe fallback action.""" available = observation.get("available_actions", []) if "advance_phase" in available: return {"action_type": "advance_phase", "rationale": "Fallback: advancing phase"} if "advance_quarter" in available: return {"action_type": "advance_quarter", "rationale": "Fallback: advancing quarter"} if "query_department" in available: return {"action_type": "query_department", "department": "Engineering", "rationale": "Fallback: scanning"} if "submit_report" in available: return {"action_type": "submit_report", "rationale": "Fallback: submitting report"} if available: return {"action_type": available[0], "rationale": "Fallback: first available action"} return {"action_type": "advance_phase", "rationale": "Fallback: no actions available"} def _summarize_quarter(self) -> None: """Generate a quarter summary using the LLM for memory compression.""" if not self.conversation_history: self.quarter_summaries.append(f"Q{self.current_quarter}: No actions recorded.") return try: messages = self.conversation_history + [ {"role": "user", "content": QUARTER_SUMMARY_PROMPT} ] response = self.client.messages.create( model=self.model, max_tokens=500, system="You are summarizing an HR management quarter. Be concise.", messages=messages, ) summary = response.content[0].text self.quarter_summaries.append(summary) except Exception as e: logger.warning(f"Summary generation failed: {e}") self.quarter_summaries.append(f"Q{self.current_quarter}: Summary unavailable.")