agents/risk_manager.py

from agents.base import Agent
from llm.prompts import RISK_MANAGER_SYSTEM, build_risk_manager_prompt


class RiskManager(Agent):
    def __init__(self, llm_client):
        super().__init__("RiskManager", RISK_MANAGER_SYSTEM, llm_client)

    def build_prompt(self, context: dict) -> str:
        return build_risk_manager_prompt(
            context.get("recommendation", {}),
            context.get("portfolio", {}),
        )

    def parse(self, raw: str) -> dict:
        result = super().parse(raw)
        action = result.get("adjusted_action", "HOLD").upper()
        if action not in ("BUY", "SELL", "HOLD"):
            action = "HOLD"
        return {
            "approved": bool(result.get("approved", True)),
            "adjusted_action": action,
            "adjusted_size": float(result.get("adjusted_size", 0.5)),
            "risk_note": str(result.get("risk_note", "")),
        }