agents/memory.py

"""
Agent Memory — stores step history for in-context retrieval
Passed back into the LLM prompt so the agent can reason over past actions.
"""

from collections import deque
from typing import Any


class AgentMemory:
    def __init__(self, maxlen: int = 20):
        self.history: deque = deque(maxlen=maxlen)

    def add(self, state: Any, action: Any, reward: float):
        self.history.append({
            "state":  state,
            "action": action,
            "reward": round(reward, 3),
        })

    def recent(self, k: int = 5) -> list:
        items = list(self.history)
        return items[-k:]

    def total_reward(self) -> float:
        return sum(h["reward"] for h in self.history)

    def __len__(self):
        return len(self.history)