kumar6591/data-quality-env

1from __future__ import annotations2 3import json4from dataclasses import dataclass5from pathlib import Path6from typing import Any7 8 9@dataclass10class MemoryItem:11    task_id: int12    seed: int13    score: float14    query_plan: list[str]15    evidence: dict[str, Any]16 17 18class MemoryStore:19    """Simple persistent memory for agent self-improvement."""20 21    def __init__(self, path: str) -> None:22        self.path = Path(path)23        self.path.parent.mkdir(parents=True, exist_ok=True)24        self._items: list[MemoryItem] = []25        self._load()26 27    def _load(self) -> None:28        if not self.path.exists():29            self._items = []30            return31        try:32            payload = json.loads(self.path.read_text())33            raw = payload.get("items", []) if isinstance(payload, dict) else []34            items: list[MemoryItem] = []35            for r in raw:36                items.append(37                    MemoryItem(38                        task_id=int(r.get("task_id", 0)),39                        seed=int(r.get("seed", 0)),40                        score=float(r.get("score", 0.0)),41                        query_plan=[str(x) for x in r.get("query_plan", [])],42                        evidence=dict(r.get("evidence", {})),43                    )44                )45            self._items = items46        except Exception:47            self._items = []48 49    def save(self) -> None:50        payload = {51            "version": 1,52            "items": [53                {54                    "task_id": i.task_id,55                    "seed": i.seed,56                    "score": i.score,57                    "query_plan": i.query_plan,58                    "evidence": i.evidence,59                }60                for i in self._items61            ],62        }63        self.path.write_text(json.dumps(payload))64 65    def add(self, item: MemoryItem, max_items: int = 500) -> None:66        self._items.append(item)67        # keep highest-scoring memories per task68        self._items.sort(key=lambda x: (x.task_id, x.score), reverse=True)69        self._items = self._items[:max_items]70 71    def top_for_task(self, task_id: int, k: int = 5) -> list[MemoryItem]:72        rows = [i for i in self._items if i.task_id == task_id]73        rows.sort(key=lambda x: x.score, reverse=True)74        return rows[:k]75 76    def query_bias(self, task_id: int, queries: list[str], k: int = 5) -> list[float]:77        """Returns additive prior bias per query from successful memories."""78        top = self.top_for_task(task_id, k=k)79        if not top:80            return [0.0 for _ in queries]81 82        bias = [0.0 for _ in queries]83        for mem in top:84            for rank, q in enumerate(mem.query_plan):85                if q in queries:86                    i = queries.index(q)87                    # Earlier query in successful run gets stronger weight.88                    bias[i] += max(0.0, 0.08 - 0.02 * rank) * max(0.0, mem.score)89        return bias90