AlgoTune/algotune_agent

"""Role-preserving, bounded command transcript from AlgoTune Agent."""
from __future__ import annotations

from ...components import Memory


class AlgoTuneMemory(Memory):
    """Keep the initial system prompt and the last five full turns per role.

    Older messages remain in chronological order but are shortened to 100 characters,
    matching ``MessageHandler._prepare_truncated_history`` in upstream AlgoTune.
    """

    def __init__(self, max_messages_per_role: int = 5):
        self.max_messages_per_role = max(1, int(max_messages_per_role))
        self.initial_user = ""
        self.transcript: list[dict[str, str]] = []

    def initialize(self, system: str, user: str = "") -> None:
        if not self.initial_user:
            # Upstream deliberately sends its initial_system_message.txt as a user message:
            # BaseInterface._load_initial_messages says this produced better responses.
            self.initial_user = f"{system}\n\n{user}".strip()

    def add(self, role: str, content: str) -> None:
        if role not in {"assistant", "user"}:
            raise ValueError(f"unsupported AlgoTune transcript role: {role!r}")
        self.transcript.append({"role": role, "content": str(content)})

    def messages(self, token_limit: int | None = None) -> list[dict[str, str]]:
        messages: list[dict[str, str]] = []
        if self.initial_user:
            messages.append({"role": "user", "content": self.initial_user})

        essential: set[int] = set()
        for role in ("user", "assistant"):
            indices = [i for i, msg in enumerate(self.transcript) if msg["role"] == role]
            essential.update(indices[-self.max_messages_per_role :])
        older_positions: list[int] = []
        for index, message in enumerate(self.transcript):
            item = dict(message)
            if index not in essential and len(item["content"]) > 100:
                item["content"] = item["content"][:100] + "..."
            if index not in essential:
                older_positions.append(len(messages))
            messages.append(item)
        # AlgoTune uses LiteLLM's tokenizer. Keep this dependency-free port conservative with the
        # standard four-characters-per-token estimate and remove only non-essential old messages.
        if token_limit and token_limit > 0:
            estimate = lambda items: sum(len(item["content"]) + 12 for item in items) // 4 + 1
            removed = 0
            while older_positions and estimate(messages) > token_limit:
                position = older_positions.pop(0) - removed
                messages.pop(position)
                removed += 1
            if removed:
                messages.insert(1, {
                    "role": "system",
                    "content": "[Older conversation history truncated due to context length limits]",
                })
        return messages

    def read(self, spec: dict | None = None) -> str:
        limit = int((spec or {}).get("top", len(self.transcript)))
        return "\n".join(
            f"{item['role']}: {item['content']}" for item in self.transcript[-limit:]
        )

    def write(self, knowledge: str, **meta) -> None:
        self.add(str(meta.get("role") or "user"), knowledge)

    def state_dict(self) -> dict:
        return {
            "initial_user": self.initial_user,
            "transcript": [dict(item) for item in self.transcript],
        }

    def load_state_dict(self, state: dict) -> None:
        if not isinstance(state, dict):
            return
        # Read the short-lived pre-port schema too, so checkpoints created during development remain
        # usable. It had split the initial text into system/user fields.
        old_system = str(state.get("initial_system") or "")
        old_user = str(state.get("initial_user") or "")
        self.initial_user = f"{old_system}\n\n{old_user}".strip() if old_system else old_user
        transcript = state.get("transcript") or []
        self.transcript = [
            {"role": str(item["role"]), "content": str(item["content"])}
            for item in transcript
            if isinstance(item, dict) and item.get("role") in {"assistant", "user"}
        ]