""" memory_core.py - CyberBrain26 Memory Layer v2 Memory types: - rule: permanent agent rules (active in prompt) - preference: user preferences (active in prompt) - project_decision: valid project decisions (active in prompt) - fact: verified technical facts (active if valid) - conversation: audit history (NOT in prompt) - test: test data (NOT in prompt) - error: diagnostics (NOT in prompt unless explicit) """ from __future__ import annotations import json import os from datetime import UTC, datetime from pathlib import Path from typing import Any try: import chromadb CHROMADB_AVAILABLE = True except ImportError: CHROMADB_AVAILABLE = False print("WARNING: ChromaDB not available") class MemoryCore: """ Core memory management with ChromaDB backend. Active kinds (included in prompt): rule, preference, project_decision, fact Inactive kinds (audit only): conversation, test, error """ ACTIVE_KINDS = {"rule", "preference", "project_decision", "fact"} INACTIVE_KINDS = {"conversation", "test", "error"} def __init__(self, db_path: str | None = None): if db_path is None: db_path = os.environ.get('CB26_DB_PATH', 'workspace/chroma_db') self.db_path = Path(db_path) self.db_path.mkdir(parents=True, exist_ok=True) self.client: Any | None = None self.collection: Any | None = None self._mock_storage: dict[str, dict] = {} if CHROMADB_AVAILABLE: self.client = chromadb.PersistentClient(path=str(self.db_path)) self.collection = self.client.get_or_create_collection( name="cyberbrain_memory", metadata={"description": "CyberBrain26 long-term memory"} ) else: print("WARNING: Running in mock mode without ChromaDB") def _generate_id(self, kind: str) -> str: """Generate unique memory ID""" timestamp = datetime.now(UTC).strftime("%Y%m%d%H%M%S%f") return f"{kind}-{timestamp}" def _metadata_to_dict( self, kind: str, status: str, source: str, confidence: str, created_at: str, supersedes: str | None, expires_at: str | None, workspace_id: str = "default", topic: str = "Unsorted", parent_topic: str = "", conversation_id: str = "", ) -> dict[str, Any]: """Convert metadata to flat dict for ChromaDB (no nested objects)""" return { "kind": kind, "status": status, "source": source, "confidence": confidence, "created_at": created_at, "supersedes": supersedes if supersedes else "", "expires_at": expires_at if expires_at else "", "workspace_id": workspace_id or "default", "topic": topic or "Unsorted", "parent_topic": parent_topic or "", "conversation_id": conversation_id or "", } def save_memory( self, content: str, kind: str = "fact", status: str = "active", source: str = "user", confidence: str = "verified", supersedes: str | None = None, expires_at: str | None = None, workspace_id: str = "default", topic: str = "Unsorted", parent_topic: str = "", conversation_id: str = "" ) -> dict[str, Any]: """ Save a new memory record. Returns: Dict with 'id', 'content', 'metadata' """ record_id = self._generate_id(kind) created_at = datetime.now(UTC).isoformat() metadata = self._metadata_to_dict( kind=kind, status=status, source=source, confidence=confidence, created_at=created_at, supersedes=supersedes, expires_at=expires_at, workspace_id=workspace_id, topic=topic, parent_topic=parent_topic, conversation_id=conversation_id ) record = {'id': record_id, 'content': content, 'metadata': metadata} if self.collection: # ChromaDB 1.5+ requires flat metadata dicts self.collection.add( documents=[content], metadatas=[metadata], ids=[record_id] ) else: # Mock storage self._mock_storage[record_id] = record # Mark superseded record if supersedes: self.update_memory(supersedes, status="superseded") return record def list_memory( self, kind: str | None = None, status: str | None = None ) -> list[dict[str, Any]]: """List memory records with optional filters""" if self.collection is not None: where: Any = None if kind or status: where = {} if kind: where["kind"] = kind if status: where["status"] = status results = self.collection.get(where=where) or {} records = [] ids = results.get("ids") or [] documents = results.get("documents") or [] metadatas = results.get("metadatas") or [] for i, doc_id in enumerate(ids): record = { "id": doc_id, "content": documents[i] if i < len(documents) else "", "metadata": metadatas[i] if i < len(metadatas) else {}, } records.append(record) return records else: # Mock storage records = list(self._mock_storage.values()) if kind: records = [r for r in records if r["metadata"]["kind"] == kind] if status: records = [r for r in records if r["metadata"]["status"] == status] return records def get_memory(self, record_id: str) -> dict[str, Any] | None: """Get a specific memory record by ID""" if self.collection is not None: results = self.collection.get(ids=[record_id]) or {} ids = results.get("ids") or [] if not ids: return None documents = results.get("documents") or [] metadatas = results.get("metadatas") or [] return { "id": ids[0], "content": documents[0] if len(documents) > 0 else "", "metadata": metadatas[0] if len(metadatas) > 0 else {}, } else: return self._mock_storage.get(record_id) def update_memory( self, record_id: str, status: str | None = None, content: str | None = None, confidence: str | None = None, expires_at: str | None = None ) -> dict[str, Any] | None: """Update an existing memory record""" existing = self.get_memory(record_id) if not existing: return None # Update metadata if status: existing['metadata']['status'] = status if confidence: existing['metadata']['confidence'] = confidence if expires_at: existing['metadata']['expires_at'] = expires_at if content: existing['content'] = content # Update in ChromaDB if self.collection: self.collection.update( ids=[record_id], documents=[existing['content']], metadatas=[existing['metadata']] ) else: self._mock_storage[record_id] = existing return existing def archive_memory(self, record_id: str) -> dict[str, Any] | None: """Archive a memory record""" return self.update_memory(record_id, status="archived") def delete_memory(self, record_id: str) -> bool: """Delete a memory record with backup""" existing = self.get_memory(record_id) if not existing: return False # Create backup backup_file = self.db_path / f"deleted_{record_id}.json" with open(backup_file, 'w') as f: json.dump(existing, f, indent=2) if self.collection: self.collection.delete(ids=[record_id]) else: del self._mock_storage[record_id] print(f"Deleted {record_id} (backup: {backup_file})") return True def retrieve_active_context( self, query: str | None = None, limit: int = 20 ) -> list[dict[str, Any]]: """ Retrieve active memory context for prompt. Filters: - Only ACTIVE_KINDS (rule, preference, project_decision, fact) - Only status='active' - Excludes superseded and expired records """ # Get all active records all_records = self.list_memory(status="active") # Filter by active kinds active_records = [ r for r in all_records if r['metadata']['kind'] in self.ACTIVE_KINDS ] # Filter out superseded records superseded_ids = { r['metadata']['supersedes'] for r in active_records if r['metadata'].get('supersedes') } active_records = [ r for r in active_records if r['id'] not in superseded_ids ] # Filter out expired records now = datetime.now(UTC).isoformat() active_records = [ r for r in active_records if not r['metadata'].get('expires_at') or r['metadata']['expires_at'] > now ] # Sort by created_at (newest first) active_records.sort( key=lambda r: r['metadata'].get('created_at', ''), reverse=True ) return active_records[:limit] def export_memory(self, output_path: str) -> int: """Export all memory records to JSON file""" all_records = self.list_memory() with open(output_path, 'w') as f: json.dump(all_records, f, indent=2) return len(all_records) def import_memory(self, input_path: str) -> int: """Import memory records from JSON file""" with open(input_path) as f: data = json.load(f) count = 0 for item in data: try: if self.collection: self.collection.add( documents=[item['content']], metadatas=[item['metadata']], ids=[item['id']] ) else: self._mock_storage[item['id']] = item count += 1 except Exception as e: print(f"Import error for {item['id']}: {e}") return count # Convenience functions _memory_instance: MemoryCore | None = None def get_memory_core() -> MemoryCore: """Get or create MemoryCore singleton""" global _memory_instance if _memory_instance is None: _memory_instance = MemoryCore() return _memory_instance def save_memory(content: str, **kwargs) -> dict[str, Any]: """Save a memory record""" return get_memory_core().save_memory(content, **kwargs) def list_memory(kind: str | None = None, status: str | None = None) -> list[dict[str, Any]]: """List memory records""" return get_memory_core().list_memory(kind=kind, status=status) def delete_memory(record_id: str) -> bool: """Delete one memory record by ID.""" return get_memory_core().delete_memory(record_id) def remember( content: str, kind: str = "fact", source: str = "manual", confidence: str = "user_confirmed", ) -> dict[str, Any]: """Explicitly save a persistent memory record.""" content = content.strip() if not content: raise ValueError("Memory content must not be empty.") return get_memory_core().save_memory( content=content, kind=kind, status="active", source=source, confidence=confidence, ) def retrieve_active_context(query: str | None = None, limit: int = 20) -> list[dict[str, Any]]: """Retrieve active memory context for prompt""" return get_memory_core().retrieve_active_context(query=query, limit=limit) if __name__ == '__main__': # Quick test print("=== MemoryCore v2 Test ===") core = MemoryCore() print("\n1. Saving test record...") record = core.save_memory( content="Test memory for CyberBrain26", kind="test", status="active" ) print(f" Saved: {record['id']}") print("\n2. Listing all records...") all_records = core.list_memory() print(f" Total: {len(all_records)} records") print("\n3. Retrieving active context...") active = core.retrieve_active_context() print(f" Active: {len(active)} records (should exclude 'test' kind)") print("\n4. Archiving test record...") core.archive_memory(record['id']) archived = core.get_memory(record['id']) print(f" Status: {archived['metadata']['status'] if archived else 'N/A'}") print("\n=== Test complete ===")