Multi_Swarm_Coevolutive/tests/unit/test_mutation_prompt_llm.py

from __future__ import annotations

import random
from dataclasses import dataclass

from multi_swarm.genome.hypothesis import HypothesisAgentGenome, ModelTier
from multi_swarm.genome.mutation_prompt_llm import (
    MUTATION_INSTRUCTIONS,
    _extract_prompt,
    is_valid_prompt,
    mutate_prompt_llm,
)

_BASE_PROMPT = (
    "Strategia mean-reversion 1h su BTC. Entry long quando RSI(14) < 30 e "
    "close > SMA(50). Exit short quando RSI(14) > 70. Stop loss 2%."
)


def _make_genome(prompt: str = _BASE_PROMPT) -> HypothesisAgentGenome:
    return HypothesisAgentGenome(
        system_prompt=prompt,
        feature_access=["close", "high", "low"],
        temperature=0.9,
        top_p=0.95,
        model_tier=ModelTier.C,
        lookback_window=200,
        cognitive_style="physicist",
    )


@dataclass
class _FakeResult:
    text: str


class _FakeLLM:
    """Mock LLMClient: ritorna una risposta configurata in input."""

    def __init__(self, response_text: str = "", raise_exc: bool = False) -> None:
        self.response_text = response_text
        self.raise_exc = raise_exc
        self.last_call: dict[str, object] | None = None

    def complete(
        self,
        genome: HypothesisAgentGenome,
        system: str,
        user: str,
        max_tokens: int = 2000,
    ) -> _FakeResult:
        self.last_call = {
            "genome_tier": genome.model_tier,
            "system": system,
            "user": user,
            "max_tokens": max_tokens,
        }
        if self.raise_exc:
            raise RuntimeError("simulated LLM failure")
        return _FakeResult(text=self.response_text)


def test_extract_prompt_from_tag() -> None:
    raw = "preambolo blah\n<prompt>Strategia RSI > 75 short, SMA(60) trend.</prompt>\nblabla"
    assert _extract_prompt(raw) == "Strategia RSI > 75 short, SMA(60) trend."


def test_extract_prompt_no_tag_returns_stripped_text() -> None:
    raw = "  Strategia momentum breakout su 1h con ATR(14)  "
    assert _extract_prompt(raw) == "Strategia momentum breakout su 1h con ATR(14)"


def test_is_valid_prompt_accepts_proper_strategy() -> None:
    new = (
        "Strategia mean-reversion 1h su BTC. Entry long quando RSI(14) < 25 e "
        "close > SMA(50). Exit short quando RSI(14) > 75."
    )
    assert is_valid_prompt(new, _BASE_PROMPT) is True


def test_is_valid_prompt_rejects_too_short() -> None:
    assert is_valid_prompt("short", _BASE_PROMPT) is False


def test_is_valid_prompt_rejects_no_strategy_keywords() -> None:
    bad = "Questo è un testo a caso che parla del meteo di domani e della pioggia."
    assert is_valid_prompt(bad, _BASE_PROMPT) is False


def test_is_valid_prompt_rejects_identical_prompt() -> None:
    assert is_valid_prompt(_BASE_PROMPT, _BASE_PROMPT) is False


def test_mutate_prompt_llm_produces_mutated_child() -> None:
    mutated = (
        "Strategia mean-reversion 1h su BTC. Entry long quando RSI(14) < 25 e "
        "close > SMA(70). Exit short quando RSI(14) > 78."
    )
    llm = _FakeLLM(response_text=f"<prompt>{mutated}</prompt>")
    parent = _make_genome()
    rng = random.Random(0)

    child = mutate_prompt_llm(parent, llm, rng)

    assert child.system_prompt == mutated
    assert child.id != parent.id
    assert child.parent_ids == [*parent.parent_ids, parent.id]
    assert child.generation == parent.generation + 1
    assert child.model_tier == ModelTier.C  # tier C preservato sul child


def test_mutate_prompt_llm_uses_mutator_tier_b_for_llm_call() -> None:
    mutated = (
        "Strategia momentum breakout 1h. Entry long quando close > SMA(60) e "
        "ATR(14) crescente. Exit con stop loss 3%."
    )
    llm = _FakeLLM(response_text=f"<prompt>{mutated}</prompt>")
    parent = _make_genome()
    rng = random.Random(0)

    mutate_prompt_llm(parent, llm, rng)

    assert llm.last_call is not None
    assert llm.last_call["genome_tier"] == ModelTier.B


def test_mutate_prompt_llm_falls_back_on_invalid_output() -> None:
    """Output troppo corto -> fallback random_mutate (cambia uno scalare)."""
    llm = _FakeLLM(response_text="<prompt>nope</prompt>")
    parent = _make_genome()
    rng = random.Random(42)

    child = mutate_prompt_llm(parent, llm, rng)

    # random_mutate preserva system_prompt, cambia uno dei 4 scalari/style.
    assert child.system_prompt == parent.system_prompt
    assert child.parent_ids == [*parent.parent_ids, parent.id]
    assert child.generation == parent.generation + 1


def test_mutate_prompt_llm_falls_back_on_identical_output() -> None:
    llm = _FakeLLM(response_text=f"<prompt>{_BASE_PROMPT}</prompt>")
    parent = _make_genome()
    rng = random.Random(42)

    child = mutate_prompt_llm(parent, llm, rng)

    assert child.system_prompt == parent.system_prompt
    assert child.generation == parent.generation + 1


def test_mutate_prompt_llm_falls_back_on_llm_exception() -> None:
    llm = _FakeLLM(raise_exc=True)
    parent = _make_genome()
    rng = random.Random(7)

    child = mutate_prompt_llm(parent, llm, rng)

    # Fallback random_mutate sempre produce un child valido.
    assert child.system_prompt == parent.system_prompt
    assert child.generation == parent.generation + 1


def test_mutate_prompt_llm_logs_mutation_cost_when_sink_provided() -> None:
    """Quando cost_tracker+repo+run_id sono forniti, la call mutator viene loggata
    con call_kind='mutation' sia in memoria sia nel repo."""
    mutated = (
        "Strategia RSI 1h evolved. Entry long quando RSI(14) < 28 e close > "
        "SMA(50). Exit short quando RSI(14) > 72."
    )

    class _R:
        text = f"<prompt>{mutated}</prompt>"
        input_tokens = 350
        output_tokens = 140

    class _FakeLLMCosted:
        def complete(self, genome, system, user, max_tokens=2000):
            return _R()

    tracker_calls = []
    repo_calls = []

    class _FakeTracker:
        def record(self, **kw):
            tracker_calls.append(kw)
            from types import SimpleNamespace
            return SimpleNamespace(cost_usd=0.0042)

    class _FakeRepo:
        def save_cost_record(self, **kw):
            repo_calls.append(kw)

    parent = _make_genome()
    child = mutate_prompt_llm(
        parent, _FakeLLMCosted(), random.Random(0),
        cost_tracker=_FakeTracker(), repo=_FakeRepo(), run_id="run-xyz",
    )
    assert child.system_prompt == mutated
    assert len(tracker_calls) == 1
    assert tracker_calls[0]["call_kind"] == "mutation"
    assert tracker_calls[0]["tier"] == ModelTier.B
    assert tracker_calls[0]["run_id"] == "run-xyz"
    assert tracker_calls[0]["agent_id"] == parent.id
    assert tracker_calls[0]["input_tokens"] == 350
    assert tracker_calls[0]["output_tokens"] == 140

    assert len(repo_calls) == 1
    assert repo_calls[0]["call_kind"] == "mutation"
    assert repo_calls[0]["tier"] == "B"
    assert repo_calls[0]["cost_usd"] == 0.0042


def test_mutate_prompt_llm_no_logging_without_sink() -> None:
    """Senza cost_tracker+repo+run_id → niente logging cost (backward compat)."""
    mutated = (
        "Strategia RSI 1h evoluta. Entry long quando RSI(14) < 25 e close > "
        "SMA(60). Exit short quando RSI(14) > 75 e ATR rising."
    )
    llm = _FakeLLM(response_text=f"<prompt>{mutated}</prompt>")
    parent = _make_genome()
    # Non solleva (anche se 0 sink forniti)
    child = mutate_prompt_llm(parent, llm, random.Random(0))
    assert child.system_prompt == mutated


def test_mutate_prompt_llm_picks_one_of_six_instructions() -> None:
    """Verifica che il system message dell'LLM includa una delle 6 istruzioni."""
    mutated = (
        "Strategia RSI 1h. Entry long quando RSI(14) < 28 e close > SMA(50). "
        "Exit short quando RSI(14) > 72."
    )
    llm = _FakeLLM(response_text=f"<prompt>{mutated}</prompt>")
    parent = _make_genome()

    mutate_prompt_llm(parent, llm, random.Random(0))

    assert llm.last_call is not None
    user_text = str(llm.last_call["user"])
    matched_keys = [k for k in MUTATION_INSTRUCTIONS if k in user_text]
    assert len(matched_keys) >= 1, f"User prompt non contiene istruzione: {user_text[:200]}"