memory/working_memory.py

"""
WorkingMemory — In-episode accumulation of facts, plans, and context.

This is the agent's "scratchpad" during a single episode. It tracks:
  - Current goal and active plan
  - Recent observations (sliding window)
  - Discovered facts and constraints
  - Pending subgoals
  - Recent errors and corrections
  - Retrieved skills from the library

Follows the Voyager concept of "working memory = prompt context":
the current observations, retrieved skills, previous code, errors,
feedback, and critique.

Spec reference: worksim_voyager_mvp_spec.md §8.3
"""
from __future__ import annotations

import copy
from collections import deque
from datetime import datetime
from typing import Any, Dict, List, Optional


class WorkingMemory:
    """In-episode working memory for the Voyager-lite agent.

    Usage:
        wm = WorkingMemory()
        wm.reset(goal="Prepare a client brief for Acme Corp")
        wm.update_observation(obs)
        wm.add_fact("Budget is $50,000")
        wm.add_subgoal("Find all relevant emails")
        snap = wm.snapshot()  # Include in LLM prompt
    """

    MAX_RECENT_OBS = 5       # Keep last N observations
    MAX_RECENT_ACTIONS = 10  # Keep last N action results
    MAX_ERRORS = 5           # Keep last N errors

    def __init__(self):
        self.reset()

    def reset(self, goal: str = "", plan: List[str] = None):
        """Reset working memory for a new episode."""
        # ── Core state ────────────────────────────────────────────────────────
        self.current_goal: str = goal
        self.active_plan: List[str] = plan or []
        self.plan_step: int = 0  # Which step of the plan we're on

        # ── Observations ──────────────────────────────────────────────────────
        self.recent_observations: deque = deque(maxlen=self.MAX_RECENT_OBS)
        self.recent_actions: deque = deque(maxlen=self.MAX_RECENT_ACTIONS)

        # ── Discovered knowledge ──────────────────────────────────────────────
        self.discovered_facts: List[Dict[str, Any]] = []
        self.discovered_constraints: List[str] = []
        self.key_entities: Dict[str, str] = {}  # entity_name -> description

        # ── Subgoals ──────────────────────────────────────────────────────────
        self.pending_subgoals: List[Dict[str, Any]] = []
        self.completed_subgoals: List[str] = []

        # ── Errors and corrections ────────────────────────────────────────────
        self.last_errors: deque = deque(maxlen=self.MAX_ERRORS)
        self.corrections_made: List[Dict[str, Any]] = []

        # ── Retrieved skills ──────────────────────────────────────────────────
        self.retrieved_skills: List[Dict[str, Any]] = []

        # ── Assets accessed ───────────────────────────────────────────────────
        self.accessed_assets: Dict[str, Dict[str, Any]] = {}  # asset_id -> summary
        self.stale_assets: List[str] = []

        # ── Metadata ──────────────────────────────────────────────────────────
        self.step_count: int = 0
        self.created_at: str = datetime.now().isoformat()

    # ── Update methods ────────────────────────────────────────────────────────

    def update_observation(self, obs: Dict[str, Any]):
        """Record a new observation from the environment."""
        self.step_count += 1
        summary = {
            "step": self.step_count,
            "time": obs.get("current_time", ""),
            "asset_count": obs.get("asset_count", 0),
            "notifications": obs.get("notifications", []),
        }
        self.recent_observations.append(summary)

        # Auto-detect new notifications as potential facts
        for notif in obs.get("notifications", []):
            if notif and notif not in [f.get("text") for f in self.discovered_facts]:
                self.add_fact(notif, source="notification")

    def update_action_result(
        self,
        tool_name: str,
        arguments: Dict[str, Any],
        result: Dict[str, Any],
        success: bool,
    ):
        """Record an action and its result."""
        entry = {
            "step": self.step_count,
            "tool_name": tool_name,
            "arguments": {k: str(v)[:100] for k, v in arguments.items()},
            "success": success,
            "status": result.get("status", "unknown"),
        }

        # Extract key info from result
        if success and "result" in result:
            res = result["result"]
            if isinstance(res, dict):
                entry["key_info"] = {k: str(v)[:200] for k, v in list(res.items())[:5]}

        self.recent_actions.append(entry)

        # Track errors
        if not success:
            self.last_errors.append({
                "step": self.step_count,
                "tool": tool_name,
                "error": result.get("error", result.get("status", "unknown")),
            })

        # Track accessed assets
        asset_id = arguments.get("asset_id", arguments.get("thread_id", arguments.get("file_id")))
        if asset_id and success:
            self.accessed_assets[asset_id] = {
                "tool": tool_name,
                "step": self.step_count,
            }

    def add_fact(self, text: str, source: str = "observation", confidence: float = 1.0):
        """Record a discovered fact."""
        # Avoid duplicates
        for f in self.discovered_facts:
            if f["text"] == text:
                return
        self.discovered_facts.append({
            "text": text,
            "source": source,
            "step": self.step_count,
            "confidence": confidence,
        })

    def add_constraint(self, constraint: str):
        """Record a discovered constraint."""
        if constraint not in self.discovered_constraints:
            self.discovered_constraints.append(constraint)

    def add_entity(self, name: str, description: str):
        """Record a key entity."""
        self.key_entities[name] = description

    def add_subgoal(self, description: str, priority: int = 1):
        """Add a pending subgoal."""
        self.pending_subgoals.append({
            "description": description,
            "priority": priority,
            "added_at_step": self.step_count,
        })

    def complete_subgoal(self, description: str):
        """Mark a subgoal as completed."""
        self.pending_subgoals = [
            sg for sg in self.pending_subgoals
            if sg["description"] != description
        ]
        self.completed_subgoals.append(description)

    def set_plan(self, plan: List[str]):
        """Set or update the active plan."""
        self.active_plan = plan
        self.plan_step = 0

    def advance_plan(self):
        """Move to the next step in the plan."""
        self.plan_step = min(self.plan_step + 1, len(self.active_plan))

    def record_correction(self, what: str, why: str):
        """Record a correction the agent made."""
        self.corrections_made.append({
            "step": self.step_count,
            "what": what,
            "why": why,
        })

    def set_retrieved_skills(self, skills: List[Dict[str, Any]]):
        """Set the currently retrieved skills from the skill library."""
        self.retrieved_skills = skills

    def mark_stale(self, asset_id: str):
        """Mark an asset as discovered-stale."""
        if asset_id not in self.stale_assets:
            self.stale_assets.append(asset_id)

    # ── Snapshot for prompting ────────────────────────────────────────────────

    def snapshot(self) -> Dict[str, Any]:
        """Create a complete snapshot of working memory for inclusion in prompts.

        Returns a dict that can be formatted into the LLM context.
        """
        return {
            "current_goal": self.current_goal,
            "active_plan": self.active_plan,
            "plan_progress": f"Step {self.plan_step}/{len(self.active_plan)}",
            "step_count": self.step_count,
            "discovered_facts": [f["text"] for f in self.discovered_facts],
            "discovered_constraints": self.discovered_constraints,
            "key_entities": self.key_entities,
            "pending_subgoals": [
                sg["description"] for sg in
                sorted(self.pending_subgoals, key=lambda x: x["priority"], reverse=True)
            ],
            "completed_subgoals": self.completed_subgoals,
            "recent_actions": list(self.recent_actions)[-5:],
            "last_errors": list(self.last_errors),
            "corrections_made": len(self.corrections_made),
            "assets_accessed": len(self.accessed_assets),
            "stale_assets": self.stale_assets,
            "retrieved_skills": [
                {"name": s.get("name", "?"), "description": s.get("description", "")}
                for s in self.retrieved_skills[:3]
            ],
        }

    def to_prompt_text(self) -> str:
        """Format working memory as a text block for the LLM prompt."""
        snap = self.snapshot()
        parts = []

        parts.append(f"GOAL: {snap['current_goal']}")

        if snap["active_plan"]:
            parts.append(f"\nPLAN ({snap['plan_progress']}):")
            for i, step in enumerate(snap["active_plan"]):
                marker = "→" if i == self.plan_step else " "
                done = "✓" if i < self.plan_step else " "
                parts.append(f"  {done}{marker} {i+1}. {step}")

        if snap["discovered_facts"]:
            parts.append(f"\nDISCOVERED FACTS ({len(snap['discovered_facts'])}):")
            for fact in snap["discovered_facts"][-5:]:
                parts.append(f"  • {fact}")

        if snap["discovered_constraints"]:
            parts.append(f"\nCONSTRAINTS:")
            for c in snap["discovered_constraints"]:
                parts.append(f"  ⚠ {c}")

        if snap["pending_subgoals"]:
            parts.append(f"\nPENDING SUBGOALS:")
            for sg in snap["pending_subgoals"]:
                parts.append(f"  □ {sg}")

        if snap["completed_subgoals"]:
            parts.append(f"\nCOMPLETED ({len(snap['completed_subgoals'])}):")
            for sg in snap["completed_subgoals"][-3:]:
                parts.append(f"  ✓ {sg}")

        if snap["last_errors"]:
            parts.append(f"\nRECENT ERRORS:")
            for err in snap["last_errors"]:
                parts.append(f"  ✗ Step {err['step']}: {err['tool']} → {err['error']}")

        if snap["recent_actions"]:
            parts.append(f"\nLAST ACTIONS:")
            for act in snap["recent_actions"]:
                s = "✓" if act["success"] else "✗"
                parts.append(f"  {s} [{act['step']}] {act['tool_name']}")

        if snap["retrieved_skills"]:
            parts.append(f"\nAVAILABLE SKILLS:")
            for sk in snap["retrieved_skills"]:
                parts.append(f"  🔧 {sk['name']}: {sk['description'][:80]}")

        parts.append(f"\nSTATS: {snap['step_count']} steps, "
                     f"{snap['assets_accessed']} assets accessed, "
                     f"{snap['corrections_made']} corrections")

        return "\n".join(parts)

    def to_compact_prompt(self, max_facts: int = 3, max_errors: int = 2) -> str:
        """Ultra-compact working memory for small models (< 2B params).

        Produces ~5-8 lines max to fit within tight token budgets.
        """
        parts = []

        # Facts (most valuable context)
        if self.discovered_facts:
            facts_text = "; ".join(
                f["text"][:60] for f in self.discovered_facts[-max_facts:]
            )
            parts.append(f"Facts: {facts_text}")

        # Current plan step (1 line)
        if self.active_plan and self.plan_step < len(self.active_plan):
            parts.append(f"Plan: step {self.plan_step+1}/{len(self.active_plan)}"
                         f" — {self.active_plan[self.plan_step][:60]}")

        # Last errors (help avoid repeating mistakes)
        if self.last_errors:
            err_text = "; ".join(
                f"{e['tool']}→{e['error'][:30]}"
                for e in list(self.last_errors)[-max_errors:]
            )
            parts.append(f"Errors: {err_text}")

        # Accessed assets count
        if self.accessed_assets:
            parts.append(f"Opened: {len(self.accessed_assets)} assets")

        return "\n".join(parts) if parts else "No context yet."