agents/memory.py

"""
LifeOS — Memory Agent (File 7 of 15)
Stores past scheduling sessions and generates actionable insights
for the Planner via Groq LLM or a deterministic fallback.

Key design points:
  - insights_cache is invalidated on every store() call
  - retrieve() returns the cache if still valid (no redundant LLM calls)
  - _generate_insights() is only called when cache is empty
  - All Groq calls use os.getenv("GROQ_API_KEY") — never hardcoded
"""
from __future__ import annotations

import json
import os
from typing import Any, Dict, List

from dotenv import load_dotenv
load_dotenv()

try:
    from groq import Groq
    _HAS_GROQ = True
except ImportError:
    _HAS_GROQ = False

from utils.prompts import MEMORY_RETRIEVAL_PROMPT


class MemoryAgent:
    """
    Manages per-episode session history and LLM-powered scheduling insights.

    Attributes
    ----------
    sessions       : list of stored session dicts
    insights_cache : cached plain-text insights string (cleared on store)
    """

    def __init__(self) -> None:
        self.sessions: List[Dict[str, Any]] = []
        self.insights_cache: str = ""

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def store(
        self,
        iteration: int,
        plan: Dict[str, Any],
        issues: List[str],
        reward: float,
    ) -> None:
        """
        Append one session record and invalidate the insights cache.
        The cache is cleared so the next retrieve() generates fresh insights
        that incorporate this new session.
        """
        self.sessions.append({
            "iteration":    iteration,
            "reward":       reward,
            "issue_count":  len(issues),
            "issues":       issues,
            "plan_summary": self._summarize_plan(plan),
        })
        self.insights_cache = ""   # force fresh insights on next retrieve()
        print(f"[Memory] Stored session {iteration}: reward={reward:.1f}, issues={len(issues)}")

    def retrieve(self) -> str:
        """
        Return insights string.
        Returns cached value if available (avoids redundant LLM calls).
        Generates fresh insights only when cache is empty.
        """
        if not self.sessions:
            return "- No history yet. Use default scheduling principles."
        if self.insights_cache:
            print("[Memory] Returning cached insights.")
            return self.insights_cache
        print("[Memory] Cache empty — generating fresh insights.")
        self.insights_cache = self._generate_insights()
        return self.insights_cache

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _generate_insights(self) -> str:
        """
        Format MEMORY_RETRIEVAL_PROMPT and call Groq.
        Falls back to _deterministic_insights() on any failure.
        """
        history_json = json.dumps(self.sessions, indent=2)
        prompt = MEMORY_RETRIEVAL_PROMPT.format(history=history_json)

        if not _HAS_GROQ:
            return self._deterministic_insights()

        api_key = os.getenv("GROQ_API_KEY")
        if not api_key:
            return self._deterministic_insights()

        try:
            client = Groq(api_key=api_key)
            response = client.chat.completions.create(
                model="llama-3.1-8b-instant",
                temperature=0.5,
                messages=[
                    {"role": "system", "content": "You are LifeOS Memory Agent. Return plain text bullet points only."},
                    {"role": "user",   "content": prompt},
                ],
            )
            result = (response.choices[0].message.content or "").strip()
            if result:
                print("[Memory] LLM insights generated successfully.")
                return result
            return self._deterministic_insights()
        except Exception as e:
            print(f"[Memory] LLM insight generation failed ({type(e).__name__}): {e}. Using deterministic fallback.")
            return self._deterministic_insights()

    def _deterministic_insights(self) -> str:
        """Produce bullet-point insights without any LLM call."""
        if not self.sessions:
            return "- No history yet. Use default scheduling principles."

        lines: List[str] = []
        rewards = [s["reward"] for s in self.sessions]
        best  = max(self.sessions, key=lambda s: s["reward"])
        worst = min(self.sessions, key=lambda s: s["reward"])

        lines.append(
            f"- Pattern: Best reward {best['reward']:.1f} at iteration {best['iteration']} "
            f"({best['plan_summary']}) -> Advice: Replicate this structure."
        )
        worst_issues = "; ".join(worst["issues"][:2]) or "none"
        lines.append(
            f"- Pattern: Worst iteration had issues: {worst_issues} "
            f"-> Advice: Directly address these in the next plan."
        )

        # Recurring issues
        all_issues: List[str] = [iss for s in self.sessions for iss in s["issues"]]
        from collections import Counter
        common = Counter(all_issues).most_common(2)
        if common:
            lines.append(
                f"- Pattern: Recurring issues: {'; '.join(i for i, _ in common)} "
                "-> Advice: Fix these first before optimizing elsewhere."
            )

        if len(rewards) >= 2:
            trend = "improving" if rewards[-1] > rewards[0] else "declining"
            lines.append(
                f"- Pattern: Reward trend is {trend} ({rewards[0]:.1f} -> {rewards[-1]:.1f}) "
                "-> Advice: Continue the current improvement strategy."
            )

        lines.append(
            "- Pattern: Always — include breaks every 90 mins and at least one revision slot "
            "-> Advice: Never omit these — they give +5 and +2 reward each."
        )
        return "\n".join(lines[:5])

    def _summarize_plan(self, plan: Dict[str, Any]) -> str:
        """Return a short human-readable summary string of a plan dict."""
        schedule = plan.get("schedule", {})
        n_days = len([d for d, s in schedule.items() if s])
        study_tasks = {
            slot.get("task")
            for slots in schedule.values()
            for slot in (slots if isinstance(slots, list) else [])
            if slot.get("type") == "study"
        }
        n_breaks = sum(
            1 for slots in schedule.values()
            for slot in (slots if isinstance(slots, list) else [])
            if slot.get("type") == "break"
        )
        return f"{len(study_tasks)} tasks, {n_days} days, {n_breaks} breaks"