task_memory.py

"""
task_memory.py — FAISS-based task memory for Phase 5 online learning.

Stores:
  - Task embeddings (512-dim float32) in FAISS IndexFlatL2
  - LoRA adapter weights on disk (one .pt file per task)
  - Task metadata (type, input snippet, score, timestamp) in JSON

Retrieval: top-k similar tasks by L2 distance → return adapters + similarity weights.
Weights: w_i = 1 / (dist_i + epsilon) — inverse distance weighting.
"""
import json, time, os
from pathlib import Path
from typing import List, Tuple, Optional

import torch
import faiss
import numpy as np

from lora import LoRAAdapter


class TaskMemory:
    DIM       = 512
    RANK      = 4
    ALPHA     = 32.0  # MUST match online_learner.LORA_ALPHA — train/inference scale must agree
    N_LAYERS  = 8

    # IndexFlatL2 returns SQUARED L2 distance. Empirically measured on the 50
    # benchmark task embeddings (phase4_latest.pt, no adapter):
    #   same-task repeat (unperturbed embed, adapter=None both times): 0.0 (squared)
    #   nearest *different* benchmark task: >=70 (squared, L2 norm >=8.4)
    # Threshold gates out irrelevant adapters (huge margin on both sides).
    # Keys are stored UNPERTURBED (embed_task(adapter=None)), matching the
    # query-time embedding exactly — see online_learner.py.
    DIST_THRESHOLD = 5.0

    # A recurring task (same canonical prompt) re-embeds to a BIT-IDENTICAL
    # vector (deterministic forward pass, eval mode, no dropout) -> dist ==
    # 0.0 exactly. Different tasks measure >=70. 1e-3 sits in the enormous
    # gap between the two, so this can never misfire on a genuinely
    # different (but nearby) task.
    DEDUP_DIST_EPS = 1e-3

    SAVE_EVERY = 10   # flush FAISS index + metadata every N adds

    def __init__(self, store_dir: str, top_k: int = 3):
        self.store_dir    = Path(store_dir)
        self.store_dir.mkdir(parents=True, exist_ok=True)
        self.top_k        = top_k
        self.index        = faiss.IndexFlatL2(self.DIM)
        self.metadata: List[dict] = []
        self._pending     = 0   # adds since last _save_index
        self._load_existing()

    # ── Persistence ──────────────────────────────────────────────────────────
    def _meta_path(self) -> Path:
        return self.store_dir / 'metadata.json'

    def _adapter_path(self, task_id: int) -> Path:
        return self.store_dir / f'adapter_{task_id:06d}.pt'

    def _index_path(self) -> Path:
        return self.store_dir / 'faiss.index'

    def _load_existing(self):
        if self._meta_path().exists():
            self.metadata = json.loads(self._meta_path().read_text())
        if self._index_path().exists() and len(self.metadata) > 0:
            self.index = faiss.read_index(str(self._index_path()))
        if len(self.metadata) > 0:
            print(f'[TaskMemory] Loaded {len(self.metadata)} tasks from {self.store_dir}')

    def _save_index(self):
        faiss.write_index(self.index, str(self._index_path()))
        self._meta_path().write_text(json.dumps(self.metadata, indent=2))

    # ── Core operations ──────────────────────────────────────────────────────
    def add(self, embedding: torch.Tensor, adapter: LoRAAdapter,
            meta: dict) -> int:
        """
        Store a task embedding + adapter. Returns task_id.
        meta: arbitrary dict (task_type, input_snippet, score, etc.)

        Dedup: if an existing entry's embedding is within DEDUP_DIST_EPS
        (squared L2) of `embedding` — i.e. the SAME recurring task — overwrite
        that entry's adapter file + metadata in place instead of appending a
        new one. The FAISS index is left untouched (the embedding is
        unchanged for a recurring task). This bounds memory size at the
        number of UNIQUE tasks ever seen, regardless of how many times each
        one recurs.
        """
        emb_np = embedding.float().cpu().numpy().reshape(1, self.DIM)
        assert emb_np.shape == (1, self.DIM)

        if self.index.ntotal > 0:
            dists, ids = self.index.search(emb_np, 1)
            if dists[0][0] < self.DEDUP_DIST_EPS:
                task_id = int(ids[0][0])
                torch.save({'state_dict': adapter.state_dict(), 'task_id': task_id},
                           self._adapter_path(task_id))
                self.metadata[task_id].update(meta)
                self.metadata[task_id]['timestamp'] = time.time()
                self._pending += 1
                if self._pending >= self.SAVE_EVERY:
                    self._save_index()
                    self._pending = 0
                return task_id

        task_id = len(self.metadata)

        # Save adapter
        torch.save({'state_dict': adapter.state_dict(), 'task_id': task_id},
                   self._adapter_path(task_id))

        # Add to FAISS
        self.index.add(emb_np)

        # Store metadata
        self.metadata.append({
            'task_id':    task_id,
            'timestamp':  time.time(),
            **meta,
        })
        self._pending += 1
        if self._pending >= self.SAVE_EVERY:
            self._save_index()
            self._pending = 0
        return task_id

    def flush(self):
        """Force-write FAISS index + metadata regardless of pending count."""
        if self._pending > 0:
            self._save_index()
            self._pending = 0

    def retrieve(self, query_emb: torch.Tensor) -> Tuple[List[LoRAAdapter], List[float]]:
        """
        Find top-k similar tasks, load their adapters.
        Returns (adapters, weights) — weights are inverse-distance normalised.
        Returns ([], []) if memory is empty.
        """
        n = self.index.ntotal
        if n == 0:
            return [], []

        k = min(self.top_k, n)
        q = query_emb.float().cpu().numpy().reshape(1, self.DIM)
        dists, ids = self.index.search(q, k)   # [1, k]
        dists = dists[0].tolist()
        ids   = ids[0].tolist()

        adapters = []
        weights  = []
        eps = 1e-6
        for dist, tid in zip(dists, ids):
            if dist > self.DIST_THRESHOLD:
                continue   # too dissimilar — irrelevant adapter would inject noise
            path = self._adapter_path(tid)
            if not path.exists():
                continue
            ckpt    = torch.load(path, map_location='cpu', weights_only=True)
            adapter = LoRAAdapter(self.N_LAYERS, self.DIM, self.RANK, self.ALPHA)
            adapter.load_state_dict(ckpt['state_dict'])
            adapters.append(adapter)
            weights.append(1.0 / (dist + eps))

        return adapters, weights

    def retrieve_merged(self, query_emb: torch.Tensor) -> Optional[LoRAAdapter]:
        """
        Retrieve top-k adapters and return a single weighted-merged adapter.
        Returns None if memory is empty.
        """
        adapters, weights = self.retrieve(query_emb)
        if not adapters:
            return None
        if len(adapters) == 1:
            return adapters[0]
        return LoRAAdapter.merged(adapters, weights,
                                  self.N_LAYERS, self.DIM, self.RANK, self.ALPHA)

    def __len__(self) -> int:
        return len(self.metadata)

    def stats(self) -> dict:
        if not self.metadata:
            return {'n_tasks': 0}
        scores = [m.get('score', 0) for m in self.metadata]
        return {
            'n_tasks':   len(self.metadata),
            'avg_score': sum(scores) / len(scores),
            'max_score': max(scores),
            'min_score': min(scores),
        }