""" Cross-session learning via ChromaDB vector store (Fix 10). LessonStore provides persistent, semantically-searchable storage for agent lessons accumulated across multiple runs. When enabled via AgentConfig.enable_cross_session_memory: 1. At iteration 0, generate_solution retrieves past lessons relevant to the current task and prepends them to the learning_log so the generator benefits from prior repair experience even on first attempt. 2. After each successful (or max-iterations) run, run_agent() persists the final learning_log entries back to the store. Design decisions: - All ChromaDB operations are wrapped in try/except — store failures never crash the agent (graceful degradation to in-memory-only behavior). - Uses PersistentClient so lessons survive process restarts. - Document IDs are SHA-256 hashes of the lesson text (truncated to 16 hex chars) enabling idempotent upserts without duplicates. - Without chromadb installed, all methods are silent no-ops. """ import hashlib import logging import os from typing import Any logger = logging.getLogger(__name__) # Semantic similarity threshold: lessons with cosine similarity above this # value are treated as paraphrases of an existing lesson and skipped. _SEMANTIC_DUP_THRESHOLD = 0.92 class LessonStore: """ Persistent vector store for agent lessons using ChromaDB. Provides semantic retrieval of past lessons to inform future repairs. Gracefully degrades to a no-op when chromadb is not installed or unavailable. """ def __init__(self, persist_dir: str = ".agent_memory") -> None: """ Initialize and connect to the ChromaDB "lessons" collection. If HF_LESSONS_REPO and HF_TOKEN environment variables are set, existing lessons are downloaded from HuggingFace and pre-populated into the collection so the agent benefits from cross-session history on startup. Args: persist_dir: Directory where ChromaDB stores its data on disk. """ self._collection = None try: import chromadb self._client = chromadb.PersistentClient(path=persist_dir) self._collection = self._client.get_or_create_collection( "lessons", metadata={"hnsw:space": "cosine"}, ) logger.info( "LessonStore initialized at '%s' — %d lessons in store", persist_dir, self.total_lessons, ) except ImportError: logger.warning( "chromadb not installed — cross-session memory disabled. " "Install with: pip install chromadb" ) except Exception as exc: logger.warning("LessonStore init failed (non-fatal): %s", exc) # Pre-populate from HuggingFace if credentials are available if self._collection is not None: repo_id = os.environ.get("HF_LESSONS_REPO", "") token = os.environ.get("HF_TOKEN", "") if repo_id and token: self._prepopulate_from_hf(repo_id, token) def store_lesson( self, lesson: str, task_category: str = "", failure_category: str = "", task_id: str = "", ) -> None: """ Persist a single lesson string with metadata for future retrieval. Uses upsert so storing the same lesson twice is idempotent. Args: lesson: The lesson text (e.g. "Always handle None inputs before indexing"). task_category: Optional category label for filtering (e.g. "text_processing"). failure_category: The failure type that produced this lesson. task_id: Short identifier for the source task. """ if self._collection is None or not lesson.strip(): return try: doc_id = hashlib.sha256(lesson.encode()).hexdigest()[:16] self._collection.upsert( ids=[doc_id], documents=[lesson], metadatas=[{ "task_category": task_category, "failure_category": failure_category, "task_id": task_id[:100], }], ) logger.debug("LessonStore: stored lesson id=%s", doc_id) except Exception as exc: logger.warning("LessonStore.store_lesson failed (non-fatal): %s", exc) def retrieve_relevant_lessons( self, query: str, task_category: str = "", n_results: int = 5, ) -> list[str]: """ Semantically retrieve the most relevant past lessons for a query. Args: query: The task description or failure message to search against. task_category: Optional filter — only return lessons from this category. n_results: Maximum number of lessons to return. Returns: List of lesson strings ranked by relevance. Empty on any failure. """ if self._collection is None or not query.strip(): return [] try: count = self.total_lessons if count == 0: return [] actual_n = min(n_results, count) kwargs: dict[str, Any] = { "query_texts": [query], "n_results": actual_n, } if task_category: kwargs["where"] = {"task_category": task_category} results = self._collection.query(**kwargs) return results.get("documents", [[]])[0] except Exception as exc: logger.warning("LessonStore.retrieve failed (non-fatal): %s", exc) return [] def _prepopulate_from_hf(self, repo_id: str, token: str) -> None: """ Download lessons from HuggingFace and upsert them into ChromaDB. Called once at init when HF_LESSONS_REPO and HF_TOKEN are set. Uses store_lesson() (idempotent upsert) so re-running never creates duplicates in the local collection. """ try: from agent.hf_memory_sync import load_lessons_from_hf hf_lessons = load_lessons_from_hf(repo_id, token) for lesson_dict in hf_lessons: self.store_lesson( lesson_dict["lesson"], task_category=lesson_dict.get("task_id", ""), failure_category=lesson_dict.get("failure_category", ""), task_id=lesson_dict.get("task_id", ""), ) if hf_lessons: logger.info( "LessonStore: pre-populated %d lessons from HF repo '%s'.", len(hf_lessons), repo_id, ) except Exception as exc: logger.warning("LessonStore: HF pre-population failed (non-fatal): %s", exc) def _is_semantic_duplicate(self, lesson: str) -> bool: """ Return True if lesson is a paraphrase of any existing ChromaDB lesson. Uses ChromaDB's cosine distance (stored as 1 - cosine_similarity). A distance < (1 - threshold) means similarity > threshold. """ if self._collection is None or self.total_lessons == 0: return False try: results = self._collection.query( query_texts=[lesson], n_results=1, include=["distances"], ) distances = results.get("distances", [[]])[0] if distances: similarity = 1.0 - distances[0] return similarity > _SEMANTIC_DUP_THRESHOLD except Exception as exc: logger.warning("LessonStore._is_semantic_duplicate failed (non-fatal): %s", exc) return False async def sync_to_hf( self, new_lessons: list[dict], router: Any = None, ) -> None: """ Push new lessons to HuggingFace with semantic + fingerprint deduplication. Deduplication layers: 1. Semantic (here): drop lessons with cosine similarity > 0.92 to any existing lesson already in the ChromaDB collection. 2. Fingerprint (in save_lessons_to_hf): drop lessons whose SHA-256 hash matches an existing JSONL entry. No-ops silently if HF_LESSONS_REPO or HF_TOKEN are not set. Args: new_lessons: Lesson dicts with keys: lesson, task_id, failure_category, timestamp, fingerprint. router: Optional LLMRouter forwarded to save_lessons_to_hf for LLM-based compaction when the file exceeds 200 entries. """ repo_id = os.environ.get("HF_LESSONS_REPO", "") token = os.environ.get("HF_TOKEN", "") if not repo_id or not token: return # Layer 1: semantic dedup against the local ChromaDB collection semantically_new = [ ld for ld in new_lessons if not self._is_semantic_duplicate(ld["lesson"]) ] if not semantically_new: logger.info("HF sync: all %d lessons are semantic duplicates — skipping.", len(new_lessons)) return logger.info( "HF sync: %d/%d lessons passed semantic dedup.", len(semantically_new), len(new_lessons), ) try: from agent.hf_memory_sync import save_lessons_to_hf await save_lessons_to_hf(semantically_new, repo_id, token, router=router) except Exception as exc: logger.warning("LessonStore.sync_to_hf failed (non-fatal): %s", exc) @property def total_lessons(self) -> int: """Return total number of lessons stored in the collection.""" if self._collection is None: return 0 try: return self._collection.count() except Exception: return 0