from __future__ import annotations import json import hashlib from pathlib import Path from typing import Any from ..evaluator import answer_format_compliant, answer_matches from ..json_utils import atomic_write_jsonl from ..task_interface import infer_output_form class RawMemory: """Bounded interaction evidence selected by failure type and diversity.""" def __init__(self, path: str | Path, *, capacity_per_task: int = 4) -> None: self.path = Path(path) self.capacity_per_task = max(int(capacity_per_task), 0) self.records = self._load() def clear(self) -> None: self.records = [] self._persist() def add( self, task_name: str, examples: list[dict[str, Any]], predictions: list[dict[str, Any]], *, anchor_ids: set[str] | None = None, replace_anchors: bool = False, ) -> None: anchor_ids = anchor_ids or set() by_id = {str(item.get("task_id")): item for item in predictions} retained = [item for item in self.records if _source(item, "task_name") != task_name] task_records = [dict(item) for item in self.records if _source(item, "task_name") == task_name] if replace_anchors: task_records = [item for item in task_records if not item.get("is_anchor")] by_task_id = {str(_source(item, "task_id")): item for item in task_records} for example in examples: task_id = str(example.get("task_id", "")) prediction = by_id.get(task_id, {}) response = prediction.get("answer") by_task_id[task_id] = { "memory_id": _memory_id(task_id), "memory_type": "raw", "internal_source": { "task_name": task_name, "task_type": str(example.get("task_type", "")), "task_id": task_id, "split": str(example.get("split", "")), }, "question": str(example.get("question", "")), "response": response, "feedback": example.get("answer"), "response_is_correct": answer_matches(response, example.get("answer"), example=example), "format_compliant": answer_format_compliant(response, example.get("answer")), "output_form": infer_output_form(str(example.get("question", "")))["kind"], "is_anchor": task_id in anchor_ids, } task_records = list(by_task_id.values()) anchors = [item for item in task_records if item.get("is_anchor")] interactions = [item for item in task_records if not item.get("is_anchor")] anchor_slots = min(len(anchors), self.capacity_per_task) interaction_slots = max(self.capacity_per_task - anchor_slots, 0) bounded = anchors[:anchor_slots] if interaction_slots: bounded.extend(_stratified_interactions(interactions, interaction_slots)) self.records = retained + bounded self._persist() def retrieval_candidates(self) -> list[dict[str, Any]]: """Expose a global, de-labelled pool; ranking belongs to the Router selector.""" return [_retrieval_view(item) for item in self.records if not item.get("is_anchor")] def _load(self) -> list[dict[str, Any]]: if not self.path.exists(): return [] records: list[dict[str, Any]] = [] with self.path.open("r", encoding="utf-8") as f: for line in f: if not line.strip(): continue try: item = json.loads(line) except json.JSONDecodeError: continue if isinstance(item, dict): records.append(item) return records def _persist(self) -> None: atomic_write_jsonl(self.path, self.records) def _retrieval_view(item: dict[str, Any]) -> dict[str, Any]: response_is_correct = item.get("response_is_correct") if not isinstance(response_is_correct, bool): response_is_correct = answer_matches(item.get("response"), item.get("feedback")) return { "memory_id": item.get("memory_id") or _memory_id(str(_source(item, "task_id"))), "memory_kind": "raw", "question": _bounded_text(item.get("question"), 1200), "previous_attempt": _bounded_text(item.get("response"), 512), "previous_attempt_correct": response_is_correct, "correct_feedback": _bounded_text(item.get("feedback"), 512), "output_form": infer_output_form(str(item.get("question", "")))["kind"], "usage_rule": ( "correct_feedback is authoritative supervision; previous_attempt is only a historical attempt " "and may be incorrect. Do not copy an incorrect attempt." ), } def _bounded_text(value: object, limit: int) -> object: if value is None or not isinstance(value, str): return value if len(value) <= limit: return value return value[:limit].rstrip() + " …[truncated]" def _source(item: dict[str, Any], key: str) -> Any: internal = item.get("internal_source") if isinstance(internal, dict) and key in internal: return internal.get(key) return item.get(key) def _memory_id(task_id: str) -> str: digest = hashlib.sha256(str(task_id).encode("utf-8")).hexdigest()[:12] return f"raw_{digest}" def _stratified_interactions( records: list[dict[str, Any]], capacity: int, ) -> list[dict[str, Any]]: """Keep format failures, answer errors, and representative successes. Buckets are exclusive. Within each bucket greedy lexical diversity avoids spending several slots on near-identical examples. Unused quota is filled from the remaining globally diverse records. """ if capacity <= 0: return [] format_failures = [item for item in records if not bool(item.get("format_compliant", True))] answer_errors = [ item for item in records if bool(item.get("format_compliant", True)) and not bool(item.get("response_is_correct", False)) ] representative = [ item for item in records if bool(item.get("format_compliant", True)) and bool(item.get("response_is_correct", False)) ] format_quota = max(round(capacity * 0.20), 1) error_quota = max(round(capacity * 0.35), 1) representative_quota = max(capacity - format_quota - error_quota, 0) selected: list[dict[str, Any]] = [] selected.extend(_diverse_take(format_failures, format_quota, selected)) selected.extend(_diverse_take(answer_errors, error_quota, selected)) selected.extend(_diverse_take(representative, representative_quota, selected)) selected_ids = {str(_source(item, "task_id")) for item in selected} remaining = [item for item in records if str(_source(item, "task_id")) not in selected_ids] selected.extend(_diverse_take(remaining, capacity - len(selected), selected)) return selected[:capacity] def _diverse_take( candidates: list[dict[str, Any]], count: int, already_selected: list[dict[str, Any]], ) -> list[dict[str, Any]]: pool = list(candidates) chosen: list[dict[str, Any]] = [] while pool and len(chosen) < count: references = already_selected + chosen if not references: best = pool[-1] else: best = max( pool, key=lambda item: ( min( 1.0 - _token_jaccard(_question_tokens(item), _question_tokens(reference)) for reference in references ), str(_source(item, "task_id")), ), ) chosen.append(best) pool.remove(best) return chosen def _question_tokens(item: dict[str, Any]) -> set[str]: return set(_normalize_text(str(item.get("question", ""))).split()) def _normalize_text(text: str) -> str: return " ".join( "".join(character.lower() if character.isalnum() else " " for character in str(text)).split() ) def _token_jaccard(left: set[str], right: set[str]) -> float: if not left or not right: return 0.0 return len(left & right) / len(left | right)