agent/hr_agent.py

"""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.")