agents/policy.py

"""
CivicAI Policy Agent

The primary decision maker outputting OpenEnv Action objects.
Can wrap an LLM, a trained RL policy, or a heuristic fallback.
"""
from __future__ import annotations
from civicai.models import SocietyState, Action, AgentMessage, SubsidyPolicy

class PolicyAgent:
    def __init__(self, name: str = "PolicyAgent", role: str = "🏛️ Lead Policymaker"):
        self.name = name
        self.role = role

    def propose_action(self, state: SocietyState, analyst_message: AgentMessage) -> tuple[Action, AgentMessage]:
        """Propose an action based on state and analyst input."""
        
        # Basic heuristic fallback policy (If no trained RL model is injected)
        tax = 0.25
        health = 0.20
        edu = 0.15
        pol = 0.10
        subsidy = SubsidyPolicy.NONE
        
        if state.inflation > 0.08:
            tax = 0.35 # Cool down economy
        elif state.employment_rate < 0.80:
            tax = 0.15 # Stimulus
            subsidy = SubsidyPolicy.INDUSTRY
            
        if state.infection_rate > 0.10:
            health = 0.30
            
        if state.crime_rate > 0.20:
            pol = 0.20
            
        action = Action(
            tax_rate=tax,
            healthcare_budget=health,
            education_budget=edu,
            police_budget=pol,
            subsidy_policy=subsidy
        )
        
        message = AgentMessage(
            agent_name=self.name,
            agent_role=self.role,
            proposal=f"Proposing Tax: {tax:.0%}, Health: {health:.0%}, Subsidies: {subsidy.value}",
            reasoning=f"Based on Analyst report: '{analyst_message.reasoning}'"
        )
        
        return action, message