agents/trader.py

from agents.base import Agent, FALLBACK_DECISION
from llm.prompts import (
    TRADER_SYSTEM,
    build_trader_prompt_A,
    build_trader_prompt_B,
    build_trader_prompt_C,
)


class Trader(Agent):
    def __init__(self, llm_client, benchmark: str = "A"):
        super().__init__("Trader", TRADER_SYSTEM, llm_client)
        self.benchmark = benchmark

    def build_prompt(self, context: dict) -> str:
        if self.benchmark == "A":
            return build_trader_prompt_A(context)
        if self.benchmark == "B":
            return build_trader_prompt_B(
                context.get("tech_analysis", {}),
                context.get("news_analysis", {}),
                context.get("portfolio", {}),
                context.get("asset", "BTC/USDT"),
                context.get("current_price", 0),
            )
        if self.benchmark == "C":
            return build_trader_prompt_C(
                context.get("research", {}),
                context.get("risk_decision", {}),
                context.get("portfolio", {}),
                context.get("asset", "BTC/USDT"),
                context.get("current_price", 0),
            )
        return build_trader_prompt_A(context)

    def parse(self, raw: str) -> dict:
        result = super().parse(raw)
        # Validate and normalize
        action = result.get("action", "HOLD").upper()
        if action not in ("BUY", "SELL", "HOLD"):
            action = "HOLD"
        size = float(result.get("size", 0.5))
        size = max(0.0, min(1.0, size))
        confidence = float(result.get("confidence", 0.5))
        confidence = max(0.0, min(1.0, confidence))
        return {
            "action": action,
            "size": size,
            "confidence": confidence,
            "reason": str(result.get("reason", "")),
        }