agent/run.py

"""CLI entry point for running the sample HR agent."""

from __future__ import annotations

import argparse
import json
import logging
import sys

from agent.hr_agent import HRAgent

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)


def run_episode(url: str, model: str = "claude-sonnet-4-6", seed: int = 42) -> float:
    """Run a complete episode against the HR environment server.

    Args:
        url: WebSocket URL of the environment server.
        model: Claude model to use.
        seed: Random seed for the environment.

    Returns:
        Final episode score.
    """
    from hr_env.client import HRProductivityEnv
    from hr_env.models import HRAction

    agent = HRAgent(model=model)

    with HRProductivityEnv(base_url=url) as env:
        result = env.reset(seed=seed)
        obs = result.observation
        logger.info(f"Episode started: {obs.message}")

        step = 0
        while not obs.done:
            # Convert observation to dict for the agent
            obs_dict = obs.model_dump()

            # Get agent's action
            action_dict = agent.decide(obs_dict)
            logger.info(
                f"Step {step} | Q{obs.current_quarter} {obs.current_phase} | "
                f"Action: {action_dict.get('action_type')} | "
                f"Rationale: {action_dict.get('rationale', 'N/A')}"
            )

            # Build and execute action
            action = HRAction(**action_dict)
            result = env.step(action)
            obs = result.observation

            if result.reward is not None:
                logger.info(f"  Reward: {result.reward}")

            step += 1

        logger.info(f"Episode complete after {step} steps. Final score: {result.reward}")
        return result.reward or 0.0


def main() -> None:
    parser = argparse.ArgumentParser(description="Run the sample HR agent")
    parser.add_argument("--url", default="ws://localhost:8000", help="Environment server WebSocket URL")
    parser.add_argument("--model", default="claude-sonnet-4-6", help="Claude model to use")
    parser.add_argument("--seed", type=int, default=42, help="Environment random seed")
    args = parser.parse_args()

    try:
        score = run_episode(args.url, args.model, args.seed)
        print(f"\nFinal Score: {score:.4f}")
    except KeyboardInterrupt:
        print("\nInterrupted.")
        sys.exit(1)
    except Exception as e:
        logger.error(f"Episode failed: {e}", exc_info=True)
        sys.exit(1)


if __name__ == "__main__":
    main()