M11: Multi-Layer Memory

# memory/buffer_memory.py
# WHAT: Sliding window buffer with optional token-aware eviction
# WHY: Naive message list grows unbounded; token awareness prevents
#      silent truncation by the model context window
# GOTCHA: Token counting uses 4-chars-per-token approximation.
#         Real tokenization (tiktoken) is more accurate but requires
#         an extra dependency. For budget management the approximation
#         is sufficient.

from __future__ import annotations
import json
from typing import Optional


class BufferMemory:
    """
    In-process message buffer with sliding window eviction.

    Eviction policy (applied in order when both limits set):
    1. If len(messages) > max_messages: drop oldest pairs
    2. If token_count() > max_tokens: drop oldest pairs until under budget
    """

    def __init__(
        self,
        max_messages: int = 20,
        max_tokens: Optional[int] = 4000,
    ) -> None:
        self.messages: list[dict] = []
        self.max_messages = max_messages
        self.max_tokens = max_tokens

    # ------------------------------------------------------------------
    # WHAT: Estimate tokens using 4-chars-per-token heuristic
    # WHY: Avoids importing a tokenizer; accurate enough for eviction
    # GOTCHA: Code-heavy messages (JSON blobs) tokenize denser than prose.
    #         If you store tool outputs in the buffer, lower max_tokens by ~15%
    # ------------------------------------------------------------------
    def _estimate_tokens(self, messages: list[dict]) -> int:
        total_chars = sum(
            len(str(m.get("content", ""))) + len(m.get("role", ""))
            for m in messages
        )
        return total_chars // 4

    def token_count(self) -> int:
        """Return estimated token count of current buffer."""
        return self._estimate_tokens(self.messages)

    # ------------------------------------------------------------------
    # WHAT: Append a message and evict if over limits
    # WHY: Enforce limits at write time so get() always returns a safe list
    # GOTCHA: Always evict in pairs (user + assistant) to preserve
    #         turn structure. Evicting a lone user message leaves the model
    #         confused about conversation flow.
    # ------------------------------------------------------------------
    def add(self, role: str, content: str) -> None:
        """Add a message and evict oldest messages if over limits."""
        self.messages.append({"role": role, "content": content})

        # Evict by count (keep pairs by evicting from front in steps of 2)
        while len(self.messages) > self.max_messages:
            self.messages.pop(0)
            if self.messages:
                self.messages.pop(0)  # evict paired turn

        # Evict by tokens
        if self.max_tokens is not None:
            while (
                len(self.messages) >= 2
                and self._estimate_tokens(self.messages) > self.max_tokens
            ):
                self.messages.pop(0)
                self.messages.pop(0)  # keep pairs

    def get(self) -> list[dict]:
        """Return current message list, safe to pass directly to the API."""
        return list(self.messages)

    def clear(self) -> None:
        """Empty the buffer (use at session end before saving episode)."""
        self.messages = []

    def __repr__(self) -> str:
        return (
            f"BufferMemory(messages={len(self.messages)}, "
            f"~{self.token_count()} tokens)"
        )


# ------------------------------------------------------------------
# WHAT: Quick smoke test — run this file directly to verify behavior
# ------------------------------------------------------------------
if __name__ == "__main__":
    buf = BufferMemory(max_messages=6, max_tokens=200)

    for i in range(5):
        buf.add("user", f"Query number {i}: what is the status of order {i}?")
        buf.add("assistant", f"Order {i} is shipped. Tracking: TRK{i:04d}.")

    print(buf)                            # should show ~3 pairs (evicted oldest)
    print(f"Messages kept: {len(buf.get())}")
    for m in buf.get():
        print(f"  [{m['role']}] {m['content'][:60]}")

# memory/vector_memory.py
# WHAT: Semantic memory layer using ChromaDB and local sentence-transformers
# WHY: Enables "find what I know about X" queries that survive process restarts
# GOTCHA: First import of SentenceTransformer downloads the model (~80 MB).
#         Subsequent runs use the local cache. Set SENTENCE_TRANSFORMERS_HOME
#         to control where the model is cached.

from __future__ import annotations
import uuid
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
from typing import Optional


class VectorMemory:
    """
    Semantic memory layer backed by ChromaDB (local SQLite) and
    sentence-transformers/all-MiniLM-L6-v2 for embeddings.

    All embeddings are computed locally — no OpenAI key required.
    """

    # WHAT: Load embedding model once at class level (shared across instances)
    # WHY: Model load takes ~1s; doing it per-instance would be wasteful
    # GOTCHA: Model is ~80 MB; first import triggers download. Subsequent
    #         runs use local cache in ~/.cache/huggingface/
    _embed_model: Optional[SentenceTransformer] = None

    @classmethod
    def _get_embedder(cls) -> SentenceTransformer:
        if cls._embed_model is None:
            cls._embed_model = SentenceTransformer("all-MiniLM-L6-v2")
        return cls._embed_model

    def __init__(
        self,
        persist_directory: str = "./chroma_memory",
        collection_name: str = "agent_memory",
        dedup_threshold: float = 0.95,
    ) -> None:
        # WHAT: Persistent client — data survives process restarts
        # GOTCHA: Use allow_reset=True only during development/testing;
        #         it wipes all collections when called
        self._client = chromadb.PersistentClient(
            path=persist_directory,
            settings=Settings(anonymized_telemetry=False),
        )
        self._collection = self._client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},  # cosine similarity index
        )
        self.dedup_threshold = dedup_threshold

    # ------------------------------------------------------------------
    # WHAT: Embed text locally and save to ChromaDB
    # WHY: Persistent storage means this memory survives agent restarts
    # GOTCHA: metadata values must be str, int, float, or bool — not dicts
    #         or nested objects. Flatten before passing.
    # ------------------------------------------------------------------
    def save(self, text: str, metadata: Optional[dict] = None) -> str:
        """
        Embed text and store in ChromaDB. Returns the generated memory ID.
        Skips saving if a near-duplicate already exists (cosine sim > dedup_threshold).
        """
        embedder = self._get_embedder()
        embedding = embedder.encode(text).tolist()

        # Deduplication check — skip if too similar to existing memory
        if self._collection.count() > 0:
            existing = self._collection.query(
                query_embeddings=[embedding],
                n_results=1,
                include=["distances"],
            )
            if existing["distances"] and existing["distances"][0]:
                # ChromaDB cosine distance = 1 - cosine_similarity
                cosine_sim = 1.0 - existing["distances"][0][0]
                if cosine_sim >= self.dedup_threshold:
                    # Near-duplicate found — skip save
                    existing_id = existing["ids"][0][0]
                    return existing_id

        memory_id = str(uuid.uuid4())
        safe_meta = {k: v for k, v in (metadata or {}).items()
                     if isinstance(v, (str, int, float, bool))}
        safe_meta["text"] = text  # store original text in metadata for recall

        self._collection.add(
            ids=[memory_id],
            embeddings=[embedding],
            metadatas=[safe_meta],
        )
        return memory_id

    # ------------------------------------------------------------------
    # WHAT: Recall top-k memories most relevant to query
    # WHY: Semantic retrieval finds related facts even with different wording
    # GOTCHA: n_results must be <= collection count; cap automatically
    # ------------------------------------------------------------------
    def recall(self, query: str, k: int = 5) -> list[dict]:
        """
        Return up to k memories most semantically similar to query.
        Each result: {"id": str, "text": str, "score": float, "metadata": dict}
        """
        count = self._collection.count()
        if count == 0:
            return []

        n = min(k, count)
        embedder = self._get_embedder()
        query_embedding = embedder.encode(query).tolist()

        results = self._collection.query(
            query_embeddings=[query_embedding],
            n_results=n,
            include=["metadatas", "distances"],
        )

        memories = []
        for i, mem_id in enumerate(results["ids"][0]):
            meta = results["metadatas"][0][i]
            distance = results["distances"][0][i]
            cosine_sim = 1.0 - distance
            memories.append({
                "id": mem_id,
                "text": meta.get("text", ""),
                "score": round(cosine_sim, 4),
                "metadata": {k: v for k, v in meta.items() if k != "text"},
            })

        return sorted(memories, key=lambda m: m["score"], reverse=True)

    # ------------------------------------------------------------------
    # WHAT: Delete a specific memory by ID
    # GOTCHA: Silently does nothing if ID not found — check return value
    #         if you need to confirm deletion
    # ------------------------------------------------------------------
    def forget(self, memory_id: str) -> bool:
        """Delete memory by ID. Returns True if found and deleted."""
        try:
            self._collection.delete(ids=[memory_id])
            return True
        except Exception:
            return False

    def count(self) -> int:
        return self._collection.count()


# ------------------------------------------------------------------
# WHAT: Smoke test — save three facts, recall with related query
# ------------------------------------------------------------------
if __name__ == "__main__":
    import tempfile, os
    with tempfile.TemporaryDirectory() as tmp:
        mem = VectorMemory(persist_directory=tmp)

        id1 = mem.save("Order TRK-001 was shipped via FedEx on Monday", {"order_id": "TRK-001"})
        id2 = mem.save("Customer prefers email notifications over SMS", {"type": "preference"})
        id3 = mem.save("The order was delayed due to a weather event in Memphis", {"order_id": "TRK-001"})

        print(f"Stored {mem.count()} memories")

        results = mem.recall("What happened with the delivery?", k=3)
        for r in results:
            print(f"  [{r['score']:.3f}] {r['text'][:70]}")

# memory/episodic_memory.py
# WHAT: Session summary memory — compresses conversations into structured episodes
# WHY: Enables multi-session continuity for long-running agents without
#      replaying full message histories
# GOTCHA: Episodes are stored as JSON files in a directory. For production,
#         store in a database or vector store for better search. This file-
#         backed approach is sufficient for <1000 sessions.

from __future__ import annotations
import json
import uuid
import re
from datetime import datetime
from pathlib import Path
from openai import OpenAI

# WHAT: Compression prompt — explicit structure prevents vague summaries
# GOTCHA: Mistral may wrap JSON in ``` fences. The save_episode() method
#         strips these before parsing. Never trust raw model output as valid JSON.
COMPRESSION_PROMPT = """You are a memory compression agent. Summarize the conversation below
into a structured episode that a future agent instance can use to quickly re-establish context.

CONVERSATION:
{conversation}

Produce a JSON object with exactly these fields:
{{
  "summary": "2-3 sentence narrative of what happened",
  "entities": ["list of named entities mentioned (people, companies, IDs, etc.)"],
  "decisions": ["list of decisions or conclusions reached"],
  "open_questions": ["list of unresolved questions or next steps"],
  "key_facts": ["list of specific facts worth remembering (numbers, dates, IDs)"]
}}

Return ONLY valid JSON — no markdown fences, no preamble."""


class EpisodicMemory:
    """
    Episodic memory layer. Compresses conversations with Mistral and stores
    structured episode files on disk. Retrieves relevant episodes by keyword
    matching on entities, decisions, and key facts.
    """

    def __init__(
        self,
        storage_dir: str = "./episodes",
        model: str = "mistral",
        base_url: str = "http://localhost:11434/v1",
        compress_after_turns: int = 20,
    ) -> None:
        self.storage_dir = Path(storage_dir)
        self.storage_dir.mkdir(parents=True, exist_ok=True)
        self.model = model
        self.compress_after_turns = compress_after_turns
        self._client = OpenAI(base_url=base_url, api_key="ollama")

    # ------------------------------------------------------------------
    # WHAT: Compress a message list into a structured episode with Mistral
    # WHY: Reduces a 50-message conversation to a ~200-token context prefix
    # GOTCHA: Mistral sometimes wraps JSON in ```json ... ``` blocks.
    #         Always strip fence markers before json.loads().
    # ------------------------------------------------------------------
    def save_episode(self, messages: list[dict], session_id: str | None = None) -> str:
        """
        Compress messages into an episode and save to disk.
        Returns the episode ID.
        """
        if not messages:
            raise ValueError("Cannot save empty episode")

        conversation = "\n".join(
            f"[{m['role'].upper()}] {m['content']}" for m in messages
        )

        try:
            resp = self._client.chat.completions.create(
                model=self.model,
                messages=[{
                    "role": "user",
                    "content": COMPRESSION_PROMPT.format(conversation=conversation),
                }],
                temperature=0,
                max_tokens=600,
            )
            raw = resp.choices[0].message.content or ""

            # WHAT: Strip markdown fences if model wraps in ``` blocks
            raw = re.sub(r"^```(?:json)?\s*", "", raw.strip())
            raw = re.sub(r"\s*```$", "", raw)

            episode_data = json.loads(raw)
        except json.JSONDecodeError:
            # Fallback: store raw summary if parsing fails
            episode_data = {
                "summary": raw[:500],
                "entities": [], "decisions": [],
                "open_questions": [], "key_facts": [],
            }

        episode_id = session_id or str(uuid.uuid4())
        episode = {
            "id": episode_id,
            "timestamp": datetime.utcnow().isoformat(),
            "turn_count": len(messages),
            **episode_data,
        }

        ep_path = self.storage_dir / f"{episode_id}.json"
        ep_path.write_text(json.dumps(episode, indent=2))
        return episode_id

    # ------------------------------------------------------------------
    # WHAT: Retrieve episodes relevant to a query by keyword overlap
    # WHY: Simple keyword matching works well for structured episode fields.
    #      For advanced use, replace with vector similarity search.
    # GOTCHA: This returns episodes ranked by overlap score, not timestamp.
    #         If you want recency bias, combine overlap with timestamp.
    # ------------------------------------------------------------------
    def recall_relevant(self, query: str, max_episodes: int = 3) -> list[dict]:
        """
        Return up to max_episodes episodes most relevant to the query.
        Relevance is scored by keyword overlap across entities, decisions, and key_facts.
        """
        query_tokens = set(query.lower().split())
        scored = []

        for ep_file in self.storage_dir.glob("*.json"):
            try:
                ep = json.loads(ep_file.read_text())
            except (json.JSONDecodeError, OSError):
                continue

            searchable = " ".join([
                ep.get("summary", ""),
                " ".join(ep.get("entities", [])),
                " ".join(ep.get("decisions", [])),
                " ".join(ep.get("key_facts", [])),
            ]).lower()

            overlap = sum(1 for t in query_tokens if t in searchable)
            if overlap > 0:
                scored.append((overlap, ep))

        scored.sort(key=lambda x: x[0], reverse=True)
        return [ep for _, ep in scored[:max_episodes]]

    def format_as_context_prefix(self, episodes: list[dict]) -> str:
        """
        Format retrieved episodes as a context prefix string to prepend to
        the system prompt or first user message.
        """
        if not episodes:
            return ""
        lines = ["[PAST SESSION CONTEXT]"]
        for ep in episodes:
            lines.append(f"Session {ep['id'][:8]} ({ep.get('timestamp','')[:10]}):")
            lines.append(f"  Summary: {ep.get('summary','')}")
            if ep.get("entities"):
                lines.append(f"  Entities: {', '.join(ep['entities'])}")
            if ep.get("open_questions"):
                lines.append(f"  Open: {'; '.join(ep['open_questions'])}")
        return "\n".join(lines)


if __name__ == "__main__":
    import tempfile
    messages = [
        {"role": "user", "content": "I need to check if ACME Logistics LLC and Acme Logistics Inc are the same company."},
        {"role": "assistant", "content": "I'll run a fuzzy match. The name similarity is 0.91. They appear to be the same entity. Recommended action: MERGE with confidence 0.87."},
        {"role": "user", "content": "What about their UCC filings in California?"},
        {"role": "assistant", "content": "I found 3 UCC filings for ACME Logistics LLC in CA, none for Acme Logistics Inc. This is consistent with a name variant, not a separate entity."},
    ]
    with tempfile.TemporaryDirectory() as tmp:
        em = EpisodicMemory(storage_dir=tmp)
        eid = em.save_episode(messages, session_id="demo-session")
        print(f"Episode saved: {eid}")
        results = em.recall_relevant("ACME entity resolution California")
        for ep in results:
            print(f"  Found: {ep['summary'][:100]}")
            print(f"  Entities: {ep['entities']}")

# memory/agent_memory.py
# WHAT: Facade over all three memory layers — single interface for agent loop
# WHY: The agent loop only needs to call build_context() and save_turn().
#      All three-layer logic is encapsulated here.
# GOTCHA: build_context() assembles in priority order:
#         1. Episodic prefix (oldest, least specific)
#         2. Vector memories injected as a synthetic [system] message
#         3. Buffer messages (most recent, highest priority)
#         This order ensures the buffer always wins over older context.

from __future__ import annotations
from .buffer_memory import BufferMemory
from .vector_memory import VectorMemory
from .episodic_memory import EpisodicMemory


class AgentMemory:
    """
    Orchestrates BufferMemory, VectorMemory, and EpisodicMemory into a single
    interface for the agent loop.
    """

    def __init__(
        self,
        buffer_max_messages: int = 20,
        buffer_max_tokens: int = 4000,
        vector_persist_dir: str = "./chroma_memory",
        episode_dir: str = "./episodes",
        vector_recall_k: int = 5,
        episode_recall_max: int = 2,
        auto_save_to_vector: bool = True,
        model: str = "mistral",
        base_url: str = "http://localhost:11434/v1",
    ) -> None:
        self.buffer = BufferMemory(
            max_messages=buffer_max_messages,
            max_tokens=buffer_max_tokens,
        )
        self.vector = VectorMemory(persist_directory=vector_persist_dir)
        self.episodic = EpisodicMemory(
            storage_dir=episode_dir, model=model, base_url=base_url
        )
        self.vector_recall_k = vector_recall_k
        self.episode_recall_max = episode_recall_max
        self.auto_save_to_vector = auto_save_to_vector
        self._turn_counter = 0

    # ------------------------------------------------------------------
    # WHAT: Assemble context from all three layers for a given query
    # WHY: Single method call gives the agent a complete, ranked context
    # GOTCHA: Vector memories are injected as a SYSTEM message so they
    #         don't look like a fake "assistant" or "user" turn, which
    #         can confuse instruction-tuned models.
    # ------------------------------------------------------------------
    def build_context(self, query: str) -> list[dict]:
        """
        Build the full context message list for a given query.

        Returns a list of messages in this order:
        1. [system] episodic context prefix (if any relevant episodes found)
        2. [system] relevant vector memories (if any found)
        3. buffer messages (recent conversation)
        """
        messages = []

        # Layer 3: Episodic prefix
        episodes = self.episodic.recall_relevant(
            query, max_episodes=self.episode_recall_max
        )
        if episodes:
            prefix = self.episodic.format_as_context_prefix(episodes)
            messages.append({"role": "system", "content": prefix})

        # Layer 2: Vector memory injection
        vector_hits = self.vector.recall(query, k=self.vector_recall_k)
        if vector_hits:
            mem_lines = ["[RELEVANT PAST FACTS]"]
            for hit in vector_hits:
                if hit["score"] >= 0.5:  # filter low-confidence hits
                    mem_lines.append(f"- [{hit['score']:.2f}] {hit['text']}")
            if len(mem_lines) > 1:
                messages.append({
                    "role": "system",
                    "content": "\n".join(mem_lines),
                })

        # Layer 1: Buffer (most recent, highest priority)
        messages.extend(self.buffer.get())
        return messages

    # ------------------------------------------------------------------
    # WHAT: Record a new turn into buffer and optionally vector store
    # WHY: Single method keeps both layers in sync on every turn
    # GOTCHA: Don't save tool outputs to the vector store — they are
    #         often too long and structurally redundant. Only save
    #         user messages and final assistant conclusions.
    # ------------------------------------------------------------------
    def save_turn(self, role: str, content: str) -> None:
        """
        Save a turn to the buffer and optionally to the vector store.
        Only user messages and assistant conclusions are vector-stored.
        """
        self.buffer.add(role, content)
        self._turn_counter += 1

        # Auto-save semantically interesting turns to vector store
        if self.auto_save_to_vector and role in ("user", "assistant"):
            self.vector.save(content, metadata={"role": role, "turn": self._turn_counter})

    def end_session(self, session_id: str | None = None) -> str | None:
        """
        Compress current buffer into an episodic summary and clear buffer.
        Returns the episode ID, or None if buffer was empty.
        """
        messages = self.buffer.get()
        if not messages:
            return None
        episode_id = self.episodic.save_episode(messages, session_id=session_id)
        self.buffer.clear()
        return episode_id


# ------------------------------------------------------------------
# WHAT: Integration test with the OpenAI-compatible agent loop
# ------------------------------------------------------------------
if __name__ == "__main__":
    from openai import OpenAI

    client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")
    memory = AgentMemory()

    SYSTEM = "You are an entity resolution assistant. Use past context if provided."

    def agent_turn(user_input: str) -> str:
        memory.save_turn("user", user_input)
        messages = [{"role": "system", "content": SYSTEM}]
        messages.extend(memory.build_context(user_input))

        resp = client.chat.completions.create(
            model="mistral",
            messages=messages,
            max_tokens=400,
            temperature=0,
        )
        answer = resp.choices[0].message.content or ""
        memory.save_turn("assistant", answer)
        return answer

    # Two turns — second should reference context from first
    print("Turn 1:", agent_turn("Is ACME Logistics LLC the same as Acme Logistics Inc?"))
    print("Turn 2:", agent_turn("What confidence score did you assign?"))

    # End session and save episode
    ep_id = memory.end_session()
    print(f"Session saved as episode: {ep_id}")

# Start Qdrant locally — persists data in ./qdrant_storage
docker run -d \
  --name qdrant \
  -p 6333:6333 \
  -p 6334:6334 \
  -v "$(pwd)/qdrant_storage:/qdrant/storage" \
  qdrant/qdrant

# REST API: http://localhost:6333
# gRPC API: localhost:6334
# Web dashboard: http://localhost:6333/dashboard

# Install Python client
pip install qdrant-client

# Verify Qdrant is running
curl http://localhost:6333/healthz
# {"title":"qdrant - vector search engine","version":"..."}

# In memory/agent_memory.py — change ONE import:

# Before (ChromaDB):
from .vector_memory import VectorMemory

# After (Qdrant):
from .vector_memory_qdrant import VectorMemoryQdrant as VectorMemory

# Everything else in AgentMemory stays identical.
# The rest of the agent loop, build_context(), and save_turn() are unchanged.