python -m venv venv
source venv/bin/activate
pip install chromadb openai tiktoken pydantic
export OPENAI_API_KEY="sk-your-key"

# memory/working_memory.py
from __future__ import annotations
from dataclasses import dataclass
import tiktoken
import time
import logging

logger = logging.getLogger(__name__)


@dataclass
class Message:
    """A single message in working memory."""
    role: str
    content: str
    timestamp: float
    tokens: int = 0


class WorkingMemory:
    """
    Sliding window of recent messages.
    
    Maintains context within token limits by evicting
    oldest messages when capacity is exceeded.
    
    Args:
        max_tokens: Maximum tokens to retain
        max_messages: Maximum number of messages
    """
    
    def __init__(self, max_tokens: int = 4000, max_messages: int = 20):
        self.max_tokens = max_tokens
        self.max_messages = max_messages
        self.messages: list[Message] = []
        self.enc = tiktoken.get_encoding("cl100k_base")

    def add(self, role: str, content: str) -> None:
        """Add a message to working memory."""
        tokens = len(self.enc.encode(content))
        msg = Message(
            role=role,
            content=content,
            timestamp=time.time(),
            tokens=tokens,
        )
        self.messages.append(msg)
        self._trim()
        logger.debug(f"Added message ({tokens} tokens), total: {self.get_total_tokens()}")

    def get_context(self) -> list[dict[str, str]]:
        """Get messages in LLM format."""
        return [{"role": m.role, "content": m.content} for m in self.messages]

    def get_total_tokens(self) -> int:
        """Get total tokens in working memory."""
        return sum(m.tokens for m in self.messages)

    def clear(self) -> None:
        """Clear all messages from working memory."""
        self.messages.clear()
        logger.debug("Working memory cleared")

    def _trim(self) -> None:
        """Remove oldest messages to stay within limits."""
        while (
            len(self.messages) > self.max_messages
            or self.get_total_tokens() > self.max_tokens
        ):
            if self.messages:
                removed = self.messages.pop(0)
                logger.debug(f"Evicted message: {removed.content[:50]}...")
            else:
                break

    def get_recent(self, n: int = 5) -> list[dict[str, str]]:
        """Get the N most recent messages."""
        return [
            {"role": m.role, "content": m.content}
            for m in self.messages[-n:]
        ]

# memory/long_term_memory.py
from __future__ import annotations
import chromadb
from openai import OpenAI
from typing import List, Dict, Optional
import time
import logging

logger = logging.getLogger(__name__)


class LongTermMemory:
    """
    Persistent vector store for all past interactions.
    
    Uses ChromaDB for semantic search over stored memories.
    Supports different content types (interactions, facts, etc.)
    """
    
    def __init__(self, collection_name: str = "agent_memory"):
        self.client = chromadb.PersistentClient(path="./memory_db")
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},
        )
        self.openai = OpenAI()
        self._id_counter = 0

    def _get_embedding(self, text: str) -> List[float]:
        """Generate embedding for text."""
        response = self.openai.embeddings.create(
            model="text-embedding-3-small",
            input=text,
        )
        return response.data[0].embedding

    def store(
        self,
        content: str,
        metadata: Optional[Dict] = None,
        content_type: str = "interaction",
    ) -> str:
        """
        Store a memory in long-term storage.
        
        Args:
            content: Text content to store
            metadata: Additional metadata
            content_type: Type of content (interaction, fact, summary)
            
        Returns:
            Memory ID
        """
        self._id_counter += 1
        memory_id = f"mem_{self._id_counter}_{int(time.time())}"
        embedding = self._get_embedding(content)
        meta = {
            "content_type": content_type,
            "timestamp": time.time(),
            "content_length": len(content),
            **(metadata or {}),
        }
        self.collection.add(
            ids=[memory_id],
            embeddings=[embedding],
            documents=[content],
            metadatas=[meta],
        )
        logger.debug(f"Stored memory: {memory_id}")
        return memory_id

    def retrieve(
        self,
        query: str,
        n_results: int = 5,
        content_type: Optional[str] = None,
    ) -> List[Dict]:
        """
        Retrieve relevant memories using semantic search.
        
        Args:
            query: Search query
            n_results: Number of results to return
            content_type: Filter by content type
            
        Returns:
            List of retrieved memories
        """
        query_embedding = self._get_embedding(query)
        kwargs = {
            "query_embeddings": [query_embedding],
            "n_results": n_results,
        }
        if content_type:
            kwargs["where"] = {"content_type": content_type}
        
        results = self.collection.query(**kwargs)
        memories = []
        for i in range(len(results["documents"][0])):
            memories.append({
                "id": results["ids"][0][i],
                "content": results["documents"][0][i],
                "metadata": results["metadatas"][0][i],
                "score": 1 - results["distances"][0][i],
            })
        return memories

    def store_interaction(self, user_msg: str, agent_msg: str) -> None:
        """Store a user-agent interaction."""
        combined = f"User: {user_msg}\nAgent: {agent_msg}"
        self.store(
            combined,
            metadata={"user_message": user_msg[:200]},
            content_type="interaction",
        )

    def store_fact(self, fact: str, source: str = "conversation") -> None:
        """Store a factual claim."""
        self.store(fact, metadata={"source": source}, content_type="fact")

    def search_facts(self, query: str, n: int = 5) -> List[str]:
        """Search for factual memories."""
        results = self.retrieve(query, n_results=n, content_type="fact")
        return [r["content"] for r in results]

# memory/summary_memory.py
from __future__ import annotations
from openai import OpenAI
from typing import List, Dict
import json
import logging

logger = logging.getLogger(__name__)

SUMMARY_PROMPT = """Summarize the following conversation, extracting:
1. Key topics discussed
2. User preferences mentioned
3. Important facts or decisions
4. Action items or follow-ups

Conversation:
{conversation}

Provide a concise summary (200-300 words):"""

FACT_EXTRACTION_PROMPT = """Extract all factual claims from this text.
Return as a JSON array of strings, each being a standalone fact.

Text: {text}

Facts (JSON array):"""


class SummaryMemory:
    """
    Compressed conversation summaries.
    
    Periodically summarizes working memory to preserve key information
    while reducing token count.
    """
    
    def __init__(self):
        self.client = OpenAI()
        self.summaries: List[str] = []
        self.extracted_facts: List[str] = []

    def summarize_conversation(self, messages: List[Dict]) -> str:
        """
        Summarize a conversation.
        
        Args:
            messages: List of message dictionaries
            
        Returns:
            Summary text
        """
        conversation = "\n".join(
            f"{m['role']}: {m['content']}" for m in messages
        )
        
        try:
            response = self.client.chat.completions.create(
                model="gpt-4-turbo-preview",
                messages=[
                    {"role": "system", "content": "You are a precise summarizer."},
                    {"role": "user", "content": SUMMARY_PROMPT.format(
                        conversation=conversation
                    )},
                ],
                temperature=0.0,
                max_tokens=500,
            )
            summary = response.choices[0].message.content
            self.summaries.append(summary)
            logger.info(f"Generated summary ({len(summary)} chars)")
            return summary
        except Exception as e:
            logger.error(f"Summarization failed: {e}")
            return ""

    def extract_facts(self, text: str) -> List[str]:
        """
        Extract factual claims from text.
        
        Args:
            text: Text to extract facts from
            
        Returns:
            List of facts
        """
        try:
            response = self.client.chat.completions.create(
                model="gpt-4-turbo-preview",
                messages=[
                    {"role": "system", "content": "Extract facts as JSON array."},
                    {"role": "user", "content": FACT_EXTRACTION_PROMPT.format(
                        text=text
                    )},
                ],
                temperature=0.0,
                max_tokens=1000,
            )
            facts = json.loads(response.choices[0].message.content)
            self.extracted_facts.extend(facts)
            logger.info(f"Extracted {len(facts)} facts")
            return facts
        except json.JSONDecodeError:
            logger.warning("Failed to parse facts as JSON")
            return []
        except Exception as e:
            logger.error(f"Fact extraction failed: {e}")
            return []

    def get_all_summaries(self) -> str:
        """Get all summaries joined."""
        return "\n\n---\n\n".join(self.summaries)

# memory/memory_manager.py
from __future__ import annotations
from memory.working_memory import WorkingMemory
from memory.long_term_memory import LongTermMemory
from memory.summary_memory import SummaryMemory
from typing import List, Dict
import logging

logger = logging.getLogger(__name__)


class MemoryManager:
    """
    Orchestrates all memory layers.
    
    Builds context for LLM by combining:
    - Working memory (recent messages)
    - Summary memory (compressed history)
    - Long-term memory (relevant past interactions)
    """
    
    def __init__(
        self,
        working_token_limit: int = 4000,
        summary_threshold: int = 20,
    ):
        self.working = WorkingMemory(max_tokens=working_token_limit)
        self.long_term = LongTermMemory()
        self.summary = SummaryMemory()
        self.summary_threshold = summary_threshold
        self.message_count = 0

    def add_user_message(self, content: str) -> None:
        """Add a user message and trigger consolidation if needed."""
        self.working.add("user", content)
        self.message_count += 1
        if self.message_count >= self.summary_threshold:
            self._consolidate()

    def add_assistant_message(self, content: str) -> None:
        """Add an assistant message."""
        self.working.add("assistant", content)

    def get_context_for_llm(self) -> List[Dict]:
        """
        Build complete context for LLM.
        
        Combines summaries, retrieved memories, and working memory
        into a single context list.
        """
        context = []
        
        # Add summaries if available
        if self.summary.summaries:
            summary_text = self.summary.get_all_summaries()
            context.append({
                "role": "system",
                "content": f"Previous conversation summaries:\n{summary_text}",
            })
        
        # Add relevant memories
        relevant_memories = self._retrieve_relevant()
        if relevant_memories:
            memory_text = "\n".join(
                f"- {m['content'][:200]}" for m in relevant_memories
            )
            context.append({
                "role": "system",
                "content": f"Relevant memories:\n{memory_text}",
            })
        
        # Add working memory
        context.extend(self.working.get_context())
        
        return context

    def search_memories(self, query: str, n: int = 5) -> List[Dict]:
        """Search long-term memories."""
        return self.long_term.retrieve(query, n_results=n)

    def store_fact(self, fact: str) -> None:
        """Store a fact in long-term memory."""
        self.long_term.store_fact(fact)

    def _retrieve_relevant(self) -> List[Dict]:
        """Retrieve memories relevant to current context."""
        if not self.working.messages:
            return []
        last_msg = self.working.messages[-1].content
        return self.long_term.retrieve(last_msg, n_results=3)

    def _consolidate(self) -> None:
        """Consolidate working memory into summaries and long-term storage."""
        messages = self.working.get_context()
        if len(messages) >= 5:
            logger.info("Consolidating memory...")
            
            # Generate summary
            self.summary.summarize_conversation(messages)
            
            # Store interactions in long-term memory
            for msg in messages:
                if msg["role"] == "user":
                    self.long_term.store_interaction(msg["content"], "See summary")
            
            # Extract and store facts
            facts = self.summary.extract_facts(
                " ".join(m["content"] for m in messages)
            )
            for fact in facts:
                self.long_term.store_fact(fact)
            
            # Clear working memory
            self.working.clear()
            self.message_count = 0
            logger.info(f"Consolidation complete: {len(facts)} facts extracted")

memory = MemoryManager()

# User mentions preferences
memory.add_user_message("I prefer Python over JavaScript")
memory.add_assistant_message("Got it, I'll focus on Python examples")

# Later, agent retrieves this memory
context = memory.get_context_for_llm()
# Includes: "User prefers Python over JavaScript"

# Store customer interactions
memory.store_fact("Customer had billing issue on 2024-01-15")
memory.store_fact("Customer prefers email communication")

# Retrieve relevant context
memories = memory.search_memories("customer billing")
# Returns relevant past interactions

# Example: Memory with privacy controls
class PrivateMemoryManager(MemoryManager):
    def store_interaction(self, user_id: str, user_msg: str, agent_msg: str):
        # Filter PII before storage
        filtered_msg = self._filter_pii(user_msg)
        super().store_interaction(filtered_msg, agent_msg)
    
    def delete_user_memories(self, user_id: str):
        # Delete all memories for a user
        self.long_term.delete_by_user(user_id)