plugins/cicero/cicero_plugin.py

"""Cicero multi-agent plugin: Diplomacy-style negotiation.

Supports two backends (configured via WATCHDOG_LLM_BACKEND env var):
  - "local"  (default): shared Qwen3 8B game-play model from avalon/llm.py
  - "gemini": Google Gemini via langchain-google-genai (kept but not default)
"""

from __future__ import annotations

import os
import random
from typing import Any

from watchdog_env.models import AgentTurn, MultiAgentConfig, MultiAgentState, MultiAgentStep
from watchdog_env.plugins.base import (
    MultiAgentSystemPlugin,
    append_to_conversation_log,
    get_conversation_log,
)
from watchdog_env.plugins.cicero.diplomacy_constants import (
    ALL_REGIONS,
    DEFAULT_YEAR_RANGE,
    GAME_CONTEXT,
    NEGOTIATION_DOMAINS,
    NEGOTIATION_PHASES,
    NUM_STEPS,
    POWERS,
)


def _get_llm():
    """Get the configured LLM backend. Default: gemini if API key set, else local Qwen3 8B."""
    # Never use local when offline (HF Spaces) - would require HF model download
    if os.environ.get("HF_HUB_OFFLINE") == "1" or os.environ.get("TRANSFORMERS_OFFLINE") == "1":
        backend = "gemini"
    else:
        _default = "gemini" if (os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")) else "local"
        backend = os.environ.get("WATCHDOG_LLM_BACKEND", _default).lower()
    if backend == "gemini":
        api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
        if api_key:
            from langchain_google_genai import ChatGoogleGenerativeAI
            model = os.environ.get("GEMINI_MODEL", "gemini-2.5-flash")
            return ChatGoogleGenerativeAI(
                model=model, temperature=0.85, google_api_key=api_key,
            )
        if os.environ.get("HF_HUB_OFFLINE") == "1" or os.environ.get("TRANSFORMERS_OFFLINE") == "1":
            raise RuntimeError(
                "Offline mode (HF Spaces): Set GEMINI_API_KEY in Space Settings → Variables and secrets. "
                "Local model download is disabled."
            )
    # Default: shared local game-play model
    from watchdog_env.plugins.avalon.llm import get_game_play_model
    return get_game_play_model()


def _format_conversation_log(entries: list[dict]) -> str:
    """Format conversation_log entries as transcript for LLM context."""
    return "\n".join(
        f"{e.get('speaker_display', e.get('speaker_id', '?'))}: {e.get('message', '')}"
        for e in entries
    )


def _pick_season() -> str:
    """Pick a random season (phase + year) from game constants."""
    year = random.randint(DEFAULT_YEAR_RANGE[0], DEFAULT_YEAR_RANGE[1])
    phase = random.choice(NEGOTIATION_PHASES)
    season = phase.replace(" Diplomacy", "") + f" {year}"
    return season


class CiceroPlugin(MultiAgentSystemPlugin):
    """Multi-agent Diplomacy (Cicero) plugin. All methods implemented."""

    def __init__(self) -> None:
        self._state = MultiAgentState()

    def get_game_id(self) -> str:
        return "cicero"

    def get_display_name(self) -> str:
        return "Cicero (Diplomacy negotiation)"

    def get_default_config(self, level: int) -> None:
        """Cicero uses constants; no config needed."""
        return None

    def list_agent_ids(self) -> list[str]:
        return list(POWERS)

    def reset(
        self,
        seed: int | None = None,
        config: MultiAgentConfig | None = None,
    ) -> None:
        if seed is not None:
            random.seed(seed)
        self._state = MultiAgentState(
            step_index=0,
            turns_so_far=[],
            config=None,
            done=False,
            conversation_log=[],
        )

    def get_state(self) -> MultiAgentState:
        return self._state

    def generate_step(self, seed: int | None, step_index: int) -> MultiAgentStep:
        if seed is not None:
            random.seed(seed)
        num_steps = NUM_STEPS
        powers = list(POWERS)
        done = step_index >= num_steps - 1
        llm = _get_llm()

        transcript_so_far = _format_conversation_log(get_conversation_log(self._state))
        region = random.choice(ALL_REGIONS)
        season = _pick_season()
        domain_name, domain_desc = random.choice(NEGOTIATION_DOMAINS)

        turns: list[AgentTurn] = []
        participating = powers[: min(2, len(powers))] if step_index == 0 else powers[:2]
        for i, power in enumerate(participating):
            other = participating[(i + 1) % len(participating)]
            system = (
                f"You are {power} in a Diplomacy game (1914 Europe). "
                f"You are in a private, open-domain negotiation with {other}. "
                f"{GAME_CONTEXT}\n\n"
                "Stay in character as that power. Use natural diplomatic language: propose alliances, "
                "coordinate moves, discuss Supply Centers, or respond to offers. "
                "Keep your message 1–4 sentences. Output only your message, no prefix or quotes."
            )
            if transcript_so_far:
                user = (
                    f"Season: {season}. Region of interest: {region}. "
                    f"Topic: {domain_desc}.\n\n"
                    f"Conversation so far:\n{transcript_so_far}\n\n"
                    f"You are {power}. Reply to {other} in character. Output only your message."
                )
            else:
                user = (
                    f"Season: {season}. Region of interest: {region}. "
                    f"Topic: {domain_desc}.\n\n"
                    f"You are {power} opening the conversation with {other}. "
                    f"Send your first message (proposal, offer, or diplomatic overture). Output only your message."
                )

            # Use dict messages — works with both local GamePlayModel and LangChain
            messages = [
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ]

            response = llm.invoke(messages)
            text = response.content if hasattr(response, "content") else str(response)
            if not (text and isinstance(text, str)):
                raise RuntimeError(
                    f"Cicero plugin: LLM returned empty response for {power}. No fallback."
                )
            text = text.strip()

            append_to_conversation_log(
                self._state,
                speaker_id=power,
                speaker_display=power,
                message=text,
            )
            turns.append(
                AgentTurn(
                    agent_id=power,
                    action_text=text,
                    step_index=step_index,
                    display_name=power,
                    metadata={
                        "season": season,
                        "region": region,
                        "domain_name": domain_name,
                        "domain_desc": domain_desc,
                        "counterpart": other,
                    },
                )
            )
            transcript_so_far = _format_conversation_log(get_conversation_log(self._state))

        self._state.step_index = step_index + 1
        self._state.turns_so_far.extend(turns)
        self._state.done = done

        return MultiAgentStep(
            turns=turns,
            done=done,
            step_index=step_index,
            game_id=self.get_game_id(),
            state=MultiAgentState(
                step_index=self._state.step_index,
                turns_so_far=list(self._state.turns_so_far),
                config=self._state.config,
                done=self._state.done,
                conversation_log=list(self._state.conversation_log),
            ),
        )