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

@dataclass
class Message:
    role: str
    content: str
    timestamp: float
    tokens: int = 0

class WorkingMemory:
    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:
        import time
        tokens = len(self.enc.encode(content))
        msg = Message(
            role=role,
            content=content,
            timestamp=time.time(),
            tokens=tokens,
        )
        self.messages.append(msg)
        self._trim()

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

    def get_total_tokens(self) -> int:
        return sum(m.tokens for m in self.messages)

    def clear(self) -> None:
        self.messages.clear()

    def _trim(self) -> None:
        while (
            len(self.messages) > self.max_messages
            or self.get_total_tokens() > self.max_tokens
        ):
            if self.messages:
                self.messages.pop(0)
            else:
                break

    def get_recent(self, n: int = 5) -> list[dict[str, str]]:
        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

class LongTermMemory:
    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]:
        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:
        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],
        )
        return memory_id

    def retrieve(
        self,
        query: str,
        n_results: int = 5,
        content_type: Optional[str] = None,
    ) -> List[Dict]:
        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:
        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:
        self.store(
            fact,
            metadata={"source": source},
            content_type="fact",
        )

    def search_facts(self, query: str, n: int = 5) -> List[str]:
        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

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:
    def __init__(self):
        self.client = OpenAI()
        self.summaries: List[str] = []
        self.extracted_facts: List[str] = []

    def summarize_conversation(self, messages: List[Dict]) -> str:
        conversation = "\n".join(
            f"{m['role']}: {m['content']}" for m in messages
        )
        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)
        return summary

    def extract_facts(self, text: str) -> List[str]:
        import json
        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,
        )
        try:
            facts = json.loads(response.choices[0].message.content)
            self.extracted_facts.extend(facts)
            return facts
        except json.JSONDecodeError:
            return []

    def get_all_summaries(self) -> str:
        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

class MemoryManager:
    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:
        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:
        self.working.add("assistant", content)

    def get_context_for_llm(self) -> List[Dict]:
        context = []
        if self.summary.summaries:
            summary_text = self.summary.get_all_summaries()
            context.append({
                "role": "system",
                "content": f"Previous conversation summaries:\n{summary_text}",
            })
        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}",
            })
        context.extend(self.working.get_context())
        return context

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

    def store_fact(self, fact: str) -> None:
        self.long_term.store_fact(fact)

    def _retrieve_relevant(self) -> List[Dict]:
        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:
        messages = self.working.get_context()
        if len(messages) >= 5:
            self.summary.summarize_conversation(messages)
            for msg in messages:
                if msg["role"] == "user":
                    self.long_term.store_interaction(
                        msg["content"],
                        "See summary",
                    )
            facts = self.summary.extract_facts(
                " ".join(m["content"] for m in messages)
            )
            for fact in facts:
                self.long_term.store_fact(fact)
            self.working.clear()
            self.message_count = 0

# Importance scoring
class ImportanceScorer:
    def __init__(self):
        from openai import OpenAI
        self.client = OpenAI()

    def score(self, text: str) -> float:
        response = self.client.chat.completions.create(
            model="gpt-4-turbo-preview",
            messages=[
                {"role": "system", "content": """Rate importance of this text 
                from 0.0 to 1.0 based on:
                - Contains user preferences or facts
                - Contains decisions or commitments
                - Contains emotional content
                - Is likely needed in future conversations
                
                Return ONLY a number."""},
                {"role": "user", "content": text},
            ],
            temperature=0.0,
            max_tokens=5,
        )
        try:
            return float(response.choices[0].message.content)
        except ValueError:
            return 0.5

import pytest
from memory.memory_manager import MemoryManager

@pytest.fixture
def memory():
    return MemoryManager(working_token_limit=1000, summary_threshold=5)

def test_working_memory(memory):
    memory.add_user_message("Hello, I like Python")
    memory.add_assistant_message("Great! I'll remember that.")
    context = memory.get_context_for_llm()
    assert len(context) >= 2

def test_long_term_storage(memory):
    memory.store_fact("User prefers dark mode")
    results = memory.search_memories("What UI preference does the user have?")
    assert len(results) > 0

def test_memory_consolidation(memory):
    for i in range(10):
        memory.add_user_message(f"Message {i}")
        memory.add_assistant_message(f"Response {i}")
    assert len(memory.summary.summaries) > 0

# main.py
from memory.memory_manager import MemoryManager
from openai import OpenAI

def main():
    memory = MemoryManager()
    llm = OpenAI()

    print("Memory-Enhanced Agent (type 'quit' to exit)")
    while True:
        user_input = input("\nYou: ").strip()
        if user_input.lower() == "quit":
            break

        memory.add_user_message(user_input)
        context = memory.get_context_for_llm()

        response = llm.chat.completions.create(
            model="gpt-4-turbo-preview",
            messages=context,
            temperature=0.7,
        )
        answer = response.choices[0].message.content
        memory.add_assistant_message(answer)
        print(f"\nAgent: {answer}")

if __name__ == "__main__":
    main()