nlp_engine.py

"""
nlp_engine.py
=============
Natural Language → Structured Scenario Extraction.
Uses the project's shared Groq client (same key/model/timeout as WildcardAgent)
to parse a user's free-text stress-test description into a machine-readable
scenario JSON consumed by script_generator.py.
"""

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

import config  # shared GROQ_API_KEY, GROQ_MODEL, GROQ_TIMEOUT

logger = logging.getLogger(__name__)

# ─────────────────────────────────────────────────────────────────── #
# SYSTEM PROMPT                                                        #
# ─────────────────────────────────────────────────────────────────── #

_SYSTEM_PROMPT = """\
You are a DeFi security expert and smart contract stress-testing engine.
Parse the natural-language simulation request and convert it into a
structured JSON scenario. You MUST respond with ONLY valid JSON — no
markdown fences, no commentary outside the JSON.

JSON schema:
{
  "title": "Short scenario title",
  "actor": "whale | retail | arb | mev | liquidator | attacker",
  "amm": {
    "id": "string",
    "category": "amm",
    "version": "1.0",
    "core_model": {
      "type": "xyk | stableswap | constant_sum | concentrated_liquidity",
      "pricing_equation": "string",
      "invariant": "string",
      "parameters": {}
    },
    "behavioral_parameters": {
      "swap_fee_percent": 0.3,
      "slippage_curve": "linear | non_linear | piecewise",
      "arbitrage_enabled": true,
      "arbitrage_reaction_speed_blocks": 1
    },
    "risk_parameters": {
      "flash_loan_sensitivity": "low | medium | high | extreme",
      "manipulation_resistance": "low | medium | high",
      "oracle_dependency": true
    }
  },
  "liquidity_pool": {
    "id": "string",
    "category": "liquidity",
    "version": "1.0",
    "core_model": {
      "initial_tvl": 100000.0,
      "liquidity_distribution": {
        "whales_percent": 50.0,
        "retail_percent": 50.0
      }
    },
    "behavioral_parameters": {
      "withdrawal_trigger_price_drop_percent": 10.0,
      "withdrawal_fraction_on_trigger": 0.5,
      "withdrawal_speed": "slow | fast | instant | exponential",
      "volatility_threshold": 5.0,
      "reentry_enabled": true,
      "reentry_delay_blocks": 10
    },
    "risk_parameters": {
      "bank_run_probability": "low | medium | high | extreme",
      "centralization_risk": "low | medium | high | extreme"
    }
  },
  "actions": [
    {
      "type": "deposit | withdraw | swap | borrow | liquidate | flash_loan | oracle_update | manipulate_price | governance_vote | drain",
      "token": "ETH | USDC | TOKEN | LP",
      "amount": <number or null>,
      "amount_unit": "tokens | eth | usd | pct",
      "speed": "instant | rapid | slow | gradual",
      "target": "pool | vault | oracle | governance | user_address"
    }
  ],
  "market_events": [
    {
      "type": "price_drop | price_spike | oracle_lag | depeg | hyperinflation | liquidity_drain | bank_run",
      "magnitude": <0.0–1.0 for drop/drain, multiplier>1 for spike>,
      "duration_steps": <int>,
      "description": "brief description"
    }
  ],
  "risk_focus": ["reentrancy","oracle_manipulation","flash_loan","liquidation_cascade","insolvency","governance"],
  "steps": <recommended simulation steps, default 50>,
  "description": "One-sentence plain English summary"
}

Rules:
- "magnitude" for a drop is 0–1 (e.g. 40% drop = 0.40).
- "magnitude" for a spike is >1 (e.g. price doubles = 2.0).
- Keep "steps" proportional to timeline (1 step ≈ 1 min of DeFi time).
- Map known historical events (FTX, Terra Luna, etc.) faithfully.
- If the text doesn't explicitly describe 'amm' or 'liquidity_pool', omit them from the JSON or use sensible defaults.
"""

_USER_PREFIX = "Parse the following stress-test scenario into JSON:\n\n"


# ─────────────────────────────────────────────────────────────────── #
# GROQ CLIENT — same init pattern as WildcardAgent                   #
# ─────────────────────────────────────────────────────────────────── #

def _init_groq_client():
    """
    Initialise Groq client using config.GROQ_API_KEY, exactly like
    WildcardAgent._init_groq().  Returns client or None.
    """
    if not config.GROQ_API_KEY:
        logger.warning("nlp_engine: GROQ_API_KEY not set — falling back to rule-based parser.")
        return None
    try:
        from groq import Groq  # type: ignore[import]
        client = Groq(api_key=config.GROQ_API_KEY)
        logger.info("nlp_engine: Groq client initialised (model=%s).", config.GROQ_MODEL)
        return client
    except ImportError:
        logger.warning("nlp_engine: 'groq' package not installed — run pip install groq.")
        return None
    except Exception as exc:
        logger.warning("nlp_engine: Groq init failed (%s) — using rule-based fallback.", exc)
        return None


# module-level singleton (created once per process)
_groq_client = _init_groq_client()


# ─────────────────────────────────────────────────────────────────── #
# RULE-BASED FALLBACK                                                  #
# ─────────────────────────────────────────────────────────────────── #

def _fallback_parse(text: str) -> Dict[str, Any]:
    """Simple keyword/regex extraction when Groq is unavailable."""
    text_l = text.lower()

    # Actor
    actor = "retail"
    if "whale" in text_l:
        actor = "whale"
    elif any(w in text_l for w in ["mev", "flash loan", "flashloan", "arbitrag"]):
        actor = "mev"
    elif "liquidator" in text_l:
        actor = "liquidator"
    elif any(w in text_l for w in ["attacker", "hack", "exploit"]):
        actor = "attacker"

    # Amount
    amount = None
    m = re.search(r'([\d,\.]+)\s*(?:million|M\b|ETH|USDC|tokens?|%)', text, re.IGNORECASE)
    if m:
        try:
            amount = float(m.group(1).replace(",", ""))
            if "million" in text_l or m.group(0).strip().upper().endswith("M"):
                amount *= 1_000_000
        except ValueError:
            pass

    # Actions
    actions = []
    if any(w in text_l for w in ["deposit", "supply", "provide liquidity"]):
        actions.append({"type": "deposit", "token": "USDC", "amount": amount, "amount_unit": "tokens", "speed": "instant", "target": "pool"})
    if any(w in text_l for w in ["withdraw", "remove liquidity", "exit"]):
        actions.append({"type": "withdraw", "token": "USDC", "amount": amount, "amount_unit": "tokens", "speed": "rapid" if "rapid" in text_l else "instant", "target": "pool"})
    if any(w in text_l for w in ["swap", "trade", "sell", "dump"]):
        actions.append({"type": "swap", "token": "ETH", "amount": amount, "amount_unit": "tokens", "speed": "instant", "target": "pool"})
    if any(w in text_l for w in ["flash loan", "flashloan"]):
        actions.append({"type": "flash_loan", "token": "ETH", "amount": amount or 10000, "amount_unit": "eth", "speed": "instant", "target": "pool"})
    if "liquidat" in text_l:
        actions.append({"type": "liquidate", "token": "USDC", "amount": None, "amount_unit": "tokens", "speed": "instant", "target": "vault"})
    if not actions:
        actions.append({"type": "swap", "token": "ETH", "amount": amount or 1000, "amount_unit": "tokens", "speed": "instant", "target": "pool"})

    # Market events
    market_events = []
    drop_m  = re.search(r'(\d+)\s*%?\s*(?:price\s+)?drop', text_l)
    spike_m = re.search(r'(\d+)\s*x\s*(?:price\s+)?(?:spike|jump|increase)', text_l)
    if drop_m:
        pct = float(drop_m.group(1)) / 100
        market_events.append({"type": "price_drop", "magnitude": pct, "duration_steps": 10, "description": f"Price drops {int(pct*100)}%"})
    if spike_m:
        mult = float(spike_m.group(1))
        market_events.append({"type": "price_spike", "magnitude": mult, "duration_steps": 5, "description": f"Price spikes {mult}x"})
    if "oracle" in text_l and any(w in text_l for w in ["lag", "manipulat", "exploit"]):
        market_events.append({"type": "oracle_lag", "magnitude": 0.1, "duration_steps": 5, "description": "Oracle price lag"})
    if any(w in text_l for w in ["insolven", "drain", "empty", "bank run"]):
        market_events.append({"type": "liquidity_drain", "magnitude": 0.9, "duration_steps": 15, "description": "Liquidity drain"})

    # Risk focus
    risk_focus = set()
    if "reentran" in text_l:                    risk_focus.add("reentrancy")
    if "oracle" in text_l:                       risk_focus.add("oracle_manipulation")
    if any(w in text_l for w in ["flash loan", "flashloan"]): risk_focus.add("flash_loan")
    if "liquidat" in text_l:                     risk_focus.add("liquidation_cascade")
    if "insolven" in text_l:                     risk_focus.add("insolvency")
    if "govern" in text_l:                       risk_focus.add("governance")
    if not risk_focus:                            risk_focus = {"insolvency"}

    return {
        "title":         "Custom Stress Test",
        "actor":         actor,
        "actions":       actions,
        "market_events": market_events,
        "risk_focus":    sorted(risk_focus),
        "steps":         50,
        "description":   text[:200],
        "_source":       "rule_based_fallback",
    }


# ─────────────────────────────────────────────────────────────────── #
# PUBLIC API                                                           #
# ─────────────────────────────────────────────────────────────────── #

def parse_scenario(
    user_text: str,
    contract_context: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Convert natural-language scenario description to structured dict.

    Uses config.GROQ_API_KEY / config.GROQ_MODEL / config.GROQ_TIMEOUT
    — the same credentials used by WildcardAgent.  Falls back to
    rule-based extraction if Groq is unavailable.

    Args:
        user_text:        Free-text scenario description.
        contract_context: Optional analyzed contract dict for richer prompts.

    Returns:
        Structured scenario dict.
    """
    # Build enriched prompt with contract context if available
    context_block = ""
    if contract_context:
        fns   = [f["name"] for f in contract_context.get("functions", [])][:10]
        vulns = [v["type"] for v in contract_context.get("vulnerabilities", [])]
        context_block = (
            f"\n\nContract context:\n"
            f"- Functions: {', '.join(fns)}\n"
            f"- Detected vulnerabilities: {', '.join(vulns) if vulns else 'None'}\n"
        )

    full_prompt = _USER_PREFIX + user_text + context_block

    if _groq_client is None:
        return _fallback_parse(user_text)

    try:
        response = _groq_client.chat.completions.create(
            model=config.GROQ_MODEL,
            messages=[
                {"role": "system", "content": _SYSTEM_PROMPT},
                {"role": "user",   "content": full_prompt},
            ],
            temperature=0.2,
            max_tokens=1024,
            timeout=config.GROQ_TIMEOUT,
        )
        raw = response.choices[0].message.content.strip()

        # Strip accidental markdown fences (same pattern as WildcardAgent)
        json_match = re.search(r'\{.*\}', raw, re.DOTALL)
        if json_match:
            scenario = json.loads(json_match.group())
            scenario["_source"] = "groq_llm"
            logger.info("nlp_engine: Groq parse OK — %s", scenario.get("title", "?"))
            return scenario

        logger.warning("nlp_engine: Groq returned non-JSON — using rule-based fallback.")

    except Exception as exc:
        logger.warning("nlp_engine: Groq call failed (%s) — using rule-based fallback.", exc)

    return _fallback_parse(user_text)