agent_config_transformer.py

"""
agent_config_transformer.py
===========================
Transforms a parsed scenario dict (from nlp_engine or scenarios.py)
into a dynamic config dictionary mapping each of the 7 Agent classes
to their specific behavioral overrides.
"""

import json
import logging
import re
from typing import Dict, Any

import config

logger = logging.getLogger(__name__)

_SYSTEM_PROMPT = """\
You are a DeFi Simulation Controller. Your job is to take a crisis scenario description
and translate it into behavioral parameter overrides for our 7 agent types.

We have exactly 7 agent classes:
1. RetailTraderAgent: 75 small accounts. (Params: trade_prob, direction_bias, capital_fraction)
2. MidTierTraderAgent: 18 active accounts. (Params: trade_prob, direction_bias)
3. CryptoWhaleAgent: 2 massive accounts. (Params: trade_prob, strategy="pump"|"dump"|"random", capital_fraction)
4. ArbitrageTraderAgent: 2 accounts hunting spread. (Params: trade_prob, arb_threshold)
5. MEVSearcherAgent: 1 bot. (Params: trade_prob, strategy="sandwich"|"frontrun"|"backrun")
6. MarketMakerAgent: 1 quoting bot. (Params: trade_prob)
7. InstitutionalParticipantAgent: 1 structured whale. (Params: trade_prob, direction_bias, capital_fraction)

Default trade probabilities are usually low (e.g., 0.05). If the scenario describes panic,
retail `trade_prob` might spike to 0.8 with `direction_bias` 0.1 (mostly selling).
If it's a whale attack, CryptoWhaleAgent `strategy` becomes "dump" with `trade_prob` 0.9.

You MUST respond with ONLY valid JSON — a dictionary where keys are exactly the agent
class names, and values are dictionaries of parameters to override. Example:
{
  "CryptoWhaleAgent": {"trade_prob": 0.9, "strategy": "dump", "capital_fraction": 0.8},
  "RetailTraderAgent": {"trade_prob": 0.6, "direction_bias": 0.1},
  "MEVSearcherAgent": {"trade_prob": 0.9, "strategy": "sandwich"}
}
Only include agents whose behavior significantly changes in the scenario.
"""


def _init_groq_client():
    if not config.GROQ_API_KEY:
        return None
    try:
        from groq import Groq
        return Groq(api_key=config.GROQ_API_KEY)
    except:
        return None

_groq_client = _init_groq_client()


def generate_agent_configs(scenario: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
    """
    Takes a scenario dictionary and uses an LLM (or fallback) to output
    the dynamic configs for the Python Mesa agents.
    """
    desc = scenario.get("description", str(scenario))
    
    if _groq_client is None:
        logger.warning("Agent Config Transformer: No Groq configured. Using primitive rule fallbacks.")
        return _fallback_configs(desc)
        
    prompt = f"Map the following scenario into agent overrides:\n{desc}"
    
    try:
        resp = _groq_client.chat.completions.create(
            model=config.GROQ_MODEL,
            messages=[
                {"role": "system", "content": _SYSTEM_PROMPT},
                {"role": "user", "content": prompt}
            ],
            temperature=0.1,
            max_tokens=512,
            timeout=config.GROQ_TIMEOUT
        )
        raw = resp.choices[0].message.content.strip()
        
        json_match = re.search(r'\{.*\}', raw, re.DOTALL)
        if json_match:
            return json.loads(json_match.group())
        return _fallback_configs(desc)
        
    except Exception as e:
        logger.error(f"Failed to generate agent configs via LLM: {e}")
        return _fallback_configs(desc)


def _fallback_configs(text: str) -> Dict[str, Dict[str, Any]]:
    """Simple regex based fallbacks if API is down."""
    text_l = text.lower()
    configs: Dict[str, Dict[str, Any]] = {}
    
    if "whale" in text_l and "dump" in text_l:
        configs["CryptoWhaleAgent"] = {"trade_prob": 0.8, "strategy": "dump", "capital_fraction": 0.9}
    elif "whale" in text_l and "pump" in text_l:
        configs["CryptoWhaleAgent"] = {"trade_prob": 0.8, "strategy": "pump", "capital_fraction": 0.9}
        
    if "sandwich" in text_l or "mev" in text_l:
        configs["MEVSearcherAgent"] = {"trade_prob": 0.9, "strategy": "sandwich"}
        
    if "panic" in text_l or "run" in text_l:
        configs["RetailTraderAgent"] = {"trade_prob": 0.8, "direction_bias": 0.1}
        
    if "arbitrage" in text_l:
        configs["ArbitrageTraderAgent"] = {"trade_prob": 0.9, "arb_threshold": 0.001}
        
    return configs