modules/cypher_memory.py

"""CYPHER V12 M4 — Persistent inter-session memory (chromadb).

CYPHER's /chat handler is currently stateless. This module adds a persistent
vector memory:
  - Categories: cve_analysis, trade_postmortem, ioc_confirmed, cybersec_insight,
    conversation, ecosystem_fact
  - store_memory(content, category, metadata) → hash_id
  - recall_memory(query, k=3, category=None) → top-k entries by similarity
  - 6 tools exposed for TOOLS_REGISTRY (M11)

Embeddings:
  - Default: chromadb's built-in sentence-transformers all-MiniLM-L6-v2 (CPU OK)
  - Fallback: chromadb default fn if st absent (still works, lower quality)
"""
from __future__ import annotations

import hashlib
import json
import logging
import time
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

VALID_CATEGORIES = (
    "cve_analysis",
    "trade_postmortem",
    "ioc_confirmed",
    "cybersec_insight",
    "conversation",
    "ecosystem_fact",
)


class CypherMemory:
    """ChromaDB-backed persistent memory for CYPHER V12.

    Single collection with `category` metadata for filtering. Hash-stable
    deterministic IDs (content+category) so duplicate stores are no-ops.
    """

    def __init__(
        self,
        persist_dir: str = "/workspace/CYPHER_V12/memory_chroma",
        collection_name: str = "cypher_memories",
        embedding_model: str = "all-MiniLM-L6-v2",
        use_st_embeddings: bool = True,
    ):
        try:
            import chromadb
            from chromadb.config import Settings
        except ImportError as e:
            raise ImportError("chromadb not installed: " + str(e))
        self._chromadb = chromadb
        Path(persist_dir).mkdir(parents=True, exist_ok=True)
        self.persist_dir = persist_dir
        self.collection_name = collection_name
        self.client = chromadb.PersistentClient(
            path=persist_dir,
            settings=Settings(anonymized_telemetry=False),
        )
        # Embedding fn: use sentence-transformers if available, else default
        self._embed_fn = None
        if use_st_embeddings:
            try:
                from chromadb.utils import embedding_functions
                self._embed_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
                    model_name=embedding_model,
                )
            except Exception as e:
                logger.warning(
                    f"sentence_transformers fallback: {type(e).__name__}: {e}"
                )
                self._embed_fn = None
        # get_or_create_collection idempotent
        if self._embed_fn is not None:
            self.collection = self.client.get_or_create_collection(
                name=collection_name,
                embedding_function=self._embed_fn,
            )
        else:
            self.collection = self.client.get_or_create_collection(
                name=collection_name,
            )

    # ───── HELPERS ─────────────────────────────────────────────

    @staticmethod
    def _hash_id(content: str, category: str) -> str:
        h = hashlib.sha256()
        h.update(category.encode("utf-8"))
        h.update(b"|")
        h.update(content.encode("utf-8"))
        return h.hexdigest()[:24]

    @staticmethod
    def _validate_category(category: str) -> str:
        if category not in VALID_CATEGORIES:
            raise ValueError(
                f"Invalid category {category!r}. Valid: {VALID_CATEGORIES}"
            )
        return category

    # ───── CORE 6 TOOLS ────────────────────────────────────────

    def store_memory(
        self,
        content: str,
        category: str,
        metadata: dict | None = None,
    ) -> str:
        self._validate_category(category)
        if not content or not content.strip():
            raise ValueError("content must be non-empty")
        mem_id = self._hash_id(content, category)
        meta = {
            "category": category,
            "ts": int(time.time()),
            "content_len": len(content),
        }
        if metadata:
            # chromadb metadata values must be scalars (str/int/float/bool)
            for k, v in metadata.items():
                if isinstance(v, (str, int, float, bool)):
                    meta[k] = v
                else:
                    meta[k] = json.dumps(v, ensure_ascii=False)[:500]
        try:
            self.collection.upsert(
                ids=[mem_id],
                documents=[content],
                metadatas=[meta],
            )
        except Exception as e:
            logger.error(f"store_memory upsert failed: {type(e).__name__}: {e}")
            raise
        return mem_id

    def recall_memory(
        self,
        query: str,
        k: int = 3,
        category: str | None = None,
    ) -> list[dict]:
        if not query or not query.strip():
            return []
        k = max(1, min(k, 50))
        kwargs: dict[str, Any] = {"query_texts": [query], "n_results": k}
        if category is not None:
            self._validate_category(category)
            kwargs["where"] = {"category": category}
        try:
            res = self.collection.query(**kwargs)
        except Exception as e:
            logger.error(f"recall_memory query failed: {type(e).__name__}: {e}")
            return []
        ids = (res.get("ids") or [[]])[0]
        docs = (res.get("documents") or [[]])[0]
        metas = (res.get("metadatas") or [[]])[0]
        distances = (res.get("distances") or [[]])[0]
        out: list[dict] = []
        for i, (mid, doc, meta, d) in enumerate(zip(ids, docs, metas, distances)):
            out.append({
                "id": mid,
                "content": doc,
                "category": (meta or {}).get("category"),
                "ts": (meta or {}).get("ts"),
                "distance": d,
                "rank": i,
                "metadata": meta,
            })
        return out

    def list_memories(
        self,
        category: str | None = None,
        limit: int = 50,
    ) -> list[dict]:
        limit = max(1, min(limit, 1000))
        kwargs: dict[str, Any] = {"limit": limit}
        if category is not None:
            self._validate_category(category)
            kwargs["where"] = {"category": category}
        try:
            res = self.collection.get(**kwargs)
        except Exception as e:
            logger.error(f"list_memories get failed: {type(e).__name__}: {e}")
            return []
        ids = res.get("ids", [])
        docs = res.get("documents", [])
        metas = res.get("metadatas", [])
        out: list[dict] = []
        for mid, doc, meta in zip(ids, docs, metas):
            out.append({
                "id": mid,
                "content": doc,
                "category": (meta or {}).get("category"),
                "ts": (meta or {}).get("ts"),
                "metadata": meta,
            })
        # Sort newest first
        out.sort(key=lambda x: x.get("ts", 0), reverse=True)
        return out

    def delete_memory(self, memory_id: str) -> bool:
        try:
            self.collection.delete(ids=[memory_id])
            return True
        except Exception as e:
            logger.error(f"delete_memory failed: {type(e).__name__}: {e}")
            return False

    def memory_stats(self) -> dict:
        try:
            total = self.collection.count()
        except Exception as e:
            logger.error(f"count failed: {type(e).__name__}: {e}")
            total = -1
        per_cat: dict[str, int] = {}
        try:
            for cat in VALID_CATEGORIES:
                res = self.collection.get(where={"category": cat}, limit=10000)
                per_cat[cat] = len(res.get("ids", []))
        except Exception as e:
            logger.warning(f"per-cat count failed: {type(e).__name__}: {e}")
        return {
            "total": total,
            "per_category": per_cat,
            "persist_dir": self.persist_dir,
            "collection": self.collection_name,
            "embedding_fn": "sentence-transformers" if self._embed_fn else "chromadb-default",
        }

    def clear_category(self, category: str) -> int:
        self._validate_category(category)
        try:
            res = self.collection.get(where={"category": category}, limit=100000)
            ids = res.get("ids", [])
            if ids:
                self.collection.delete(ids=ids)
            return len(ids)
        except Exception as e:
            logger.error(f"clear_category failed: {type(e).__name__}: {e}")
            return 0


__all__ = ["CypherMemory", "VALID_CATEGORIES"]


if __name__ == "__main__":
    import shutil
    logging.basicConfig(level=logging.INFO)
    print("=== M4 cypher_memory SMOKE TEST ===")

    # Use clean smoke dir to avoid polluting prod
    smoke_dir = "/tmp/smoke_chroma_cypher"
    if Path(smoke_dir).exists():
        shutil.rmtree(smoke_dir, ignore_errors=True)

    mem = CypherMemory(
        persist_dir=smoke_dir,
        collection_name="smoke_test",
        use_st_embeddings=True,
    )
    print(f"Init OK. Embedding fn: {'st' if mem._embed_fn else 'chromadb-default'}")

    # Store 5 memories across categories
    samples = [
        ("CVE-2021-44228 Log4Shell critical RCE in Apache Log4j2", "cve_analysis"),
        ("Trade BTCUSDT long 42000 entry, SL 41500, TP 43500, won +3.2%", "trade_postmortem"),
        ("IOC: domain evil-c2.xyz confirmed malicious by Jescy", "ioc_confirmed"),
        ("Pattern: PowerShell -enc base64 often precedes lateral movement", "cybersec_insight"),
        ("NEXUS is the Python coding ASI, ARCHON is filesystem-trained", "ecosystem_fact"),
    ]
    ids: list[str] = []
    for content, cat in samples:
        mid = mem.store_memory(content, cat, metadata={"smoke": True})
        ids.append(mid)
        print(f"  stored id={mid} cat={cat}")

    # Recall
    res = mem.recall_memory("What is Log4Shell?", k=3)
    print(f"\nRecall 'Log4Shell': {len(res)} hits")
    for r in res:
        print(f"  rank={r['rank']} cat={r['category']} dist={r['distance']:.3f} content={r['content'][:80]}")
    assert res[0]["category"] == "cve_analysis", "Log4Shell should match cve_analysis"

    # Recall with category filter
    res2 = mem.recall_memory("BTC long position", k=2, category="trade_postmortem")
    print(f"\nFiltered recall 'BTC long' in trade_postmortem: {len(res2)} hits")
    assert len(res2) >= 1 and res2[0]["category"] == "trade_postmortem"

    # List
    listed = mem.list_memories(category="cve_analysis", limit=10)
    print(f"\nList cve_analysis: {len(listed)} entries")

    # Stats
    stats = mem.memory_stats()
    print(f"\nStats: total={stats['total']} per_cat={stats['per_category']} embed={stats['embedding_fn']}")

    # Delete
    deleted = mem.delete_memory(ids[0])
    print(f"\nDelete id[0]: {deleted}")
    assert deleted

    # Clear category
    n_cleared = mem.clear_category("trade_postmortem")
    print(f"Cleared trade_postmortem: {n_cleared}")

    print("=== SMOKE PASS ===")