Implementing Memory in Our Agent: Building a Complete Personal Assistant with Short-Term and Long-Term Memory

class ConversationManager:
    """
    Manages short-term conversation memory with automatic windowing.
    """
    
    def __init__(self, max_messages: int = 20):
        self.messages = []
        self.max_messages = max_messages
    
    def add_user_message(self, content: str):
        """Add a user message to history."""
        self.messages.append({
            "role": "user",
            "content": content
        })
        self._trim_if_needed()
    
    def add_assistant_message(self, content: str):
        """Add an assistant message to history."""
        self.messages.append({
            "role": "assistant",
            "content": content
        })
        self._trim_if_needed()
    
    def _trim_if_needed(self):
        """Keep only the most recent messages."""
        if len(self.messages) > self.max_messages:
            self.messages = self.messages[-self.max_messages:]
    
    def get_messages(self) -> list:
        """Get all current messages."""
        return self.messages.copy()
    
    def clear(self):
        """Start a fresh conversation."""
        self.messages = []

import json
import os
from datetime import datetime

class KnowledgeStore:
    """
    Stores and retrieves long-term facts about the user.
    """
    
    def __init__(self, storage_file: str = "knowledge.json"):
        self.storage_file = storage_file
        self.facts = self._load_facts()
    
    def _load_facts(self) -> list:
        """Load facts from disk."""
        if os.path.exists(self.storage_file):
            with open(self.storage_file, 'r') as f:
                return json.load(f)
        return []
    
    def _save_facts(self):
        """Save facts to disk."""
        with open(self.storage_file, 'w') as f:
            json.dump(self.facts, f, indent=2)
    
    def add_fact(self, fact: str, category: str = None):
        """Store a new fact."""
        entry = {
            "fact": fact,
            "category": category,
            "timestamp": datetime.now().isoformat()
        }
        self.facts.append(entry)
        self._save_facts()
    
    def search(self, query: str) -> list:
        """
        Search for relevant facts.
        Simple keyword matching for now.
        """
        query_lower = query.lower()
        results = []
        
        for entry in self.facts:
            fact_lower = entry["fact"].lower()
            if query_lower in fact_lower:
                results.append(entry["fact"])
            elif entry["category"] and query_lower in entry["category"].lower():
                results.append(entry["fact"])
        
        return results
    
    def get_all_facts(self) -> list:
        """Get all stored facts."""
        return [entry["fact"] for entry in self.facts]

import os
from anthropic import Anthropic

class PersonalAssistant:
    """
    A personal assistant with both short-term and long-term memory.
    """
    
    def __init__(self, api_key: str):
        # Using Claude Sonnet 4.5 for its excellent agent reasoning capabilities
        self.client = Anthropic(api_key=api_key)
        self.model = "claude-sonnet-4-5"
        
        # Memory systems
        self.conversation = ConversationManager(max_messages=20)
        self.knowledge = KnowledgeStore()
    
    def chat(self, user_message: str) -> str:
        """
        Process a user message and return a response.
        Handles both conversation and knowledge storage.
        """
        # Check if this is a memory command
        if self._is_memory_command(user_message):
            return self._handle_memory_command(user_message)
        
        # Search for relevant facts
        relevant_facts = self.knowledge.search(user_message)
        
        # Build system prompt with context
        system_prompt = self._build_system_prompt(relevant_facts)
        
        # Add user message to conversation
        self.conversation.add_user_message(user_message)
        
        # Get response from Claude
        response = self.client.messages.create(
            model=self.model,
            max_tokens=1024,
            system=system_prompt,
            messages=self.conversation.get_messages()
        )
        
        # Extract and store response
        assistant_message = response.content[0].text
        self.conversation.add_assistant_message(assistant_message)
        
        return assistant_message
    
    def _is_memory_command(self, message: str) -> bool:
        """Detect if user wants to store something."""
        keywords = ["remember", "store", "save", "keep in mind", "note that"]
        message_lower = message.lower()
        return any(keyword in message_lower for keyword in keywords)
    
    def _handle_memory_command(self, message: str) -> str:
        """Extract and store information."""
        # Use Claude to extract the fact
        response = self.client.messages.create(
            model=self.model,
            max_tokens=256,
            system="Extract the key fact the user wants you to remember. Return only the fact as a clear, concise statement.",
            messages=[{"role": "user", "content": message}]
        )
        
        fact = response.content[0].text
        self.knowledge.add_fact(fact)
        
        return f"Got it! I'll remember that {fact.lower()}"
    
    def _build_system_prompt(self, relevant_facts: list) -> str:
        """Build system prompt with retrieved knowledge."""
        base_prompt = "You are a helpful personal assistant."
        
        if not relevant_facts:
            return base_prompt
        
        facts_text = "\n".join(f"- {fact}" for fact in relevant_facts)
        return f"""{base_prompt}

You have access to the following information about the user:

{facts_text}

Use this information when relevant to provide personalized responses."""
    
    def start_new_conversation(self):
        """Clear conversation history but keep long-term knowledge."""
        self.conversation.clear()

## Create the assistant
assistant = PersonalAssistant(api_key=os.getenv("ANTHROPIC_API_KEY"))

## Store some facts
print(assistant.chat("Remember that my birthday is July 20"))
## Output: Got it! I'll remember that your birthday is july 20

print(assistant.chat("Remember that I'm allergic to peanuts"))
## Output: Got it! I'll remember that you're allergic to peanuts

print(assistant.chat("Remember that I prefer Italian food"))
## Output: Got it! I'll remember that you prefer italian food

## Now ask questions that use this knowledge
print(assistant.chat("What should I be careful about when eating out?"))
## Output: Based on what I know, you should be careful about peanuts since 
## you're allergic to them. Always inform restaurant staff about your allergy 
## and ask about ingredients, especially in sauces and desserts where peanuts 
## might be hidden...

print(assistant.chat("Suggest a restaurant for my birthday dinner"))
## Output: For your birthday on July 20, I'd suggest an Italian restaurant 
## since you prefer Italian food. Make sure to mention your peanut allergy 
## when making the reservation...

from sentence_transformers import SentenceTransformer
import numpy as np

class SemanticKnowledgeStore:
    """
    Knowledge store with semantic search using embeddings.
    """
    
    def __init__(self, storage_file: str = "knowledge.json"):
        self.storage_file = storage_file
        # Using a local embedding model to keep costs down
        self.encoder = SentenceTransformer('all-MiniLM-L6-v2')
        self.facts = []
        self.embeddings = []
        self._load_data()
    
    def _load_data(self):
        """Load facts and regenerate embeddings."""
        if os.path.exists(self.storage_file):
            with open(self.storage_file, 'r') as f:
                data = json.load(f)
                self.facts = data
                # Regenerate embeddings
                for entry in self.facts:
                    embedding = self.encoder.encode(entry["fact"])
                    self.embeddings.append(embedding)
    
    def _save_data(self):
        """Save facts to disk."""
        with open(self.storage_file, 'w') as f:
            json.dump(self.facts, f, indent=2)
    
    def add_fact(self, fact: str, category: str = None):
        """Store a fact with its embedding."""
        entry = {
            "fact": fact,
            "category": category,
            "timestamp": datetime.now().isoformat()
        }
        
        # Generate embedding
        embedding = self.encoder.encode(fact)
        
        self.facts.append(entry)
        self.embeddings.append(embedding)
        self._save_data()
    
    def search(self, query: str, top_k: int = 3, threshold: float = 0.3) -> list:
        """
        Search for facts semantically similar to the query.
        
        Args:
            query: The search query
            top_k: Maximum number of results to return
            threshold: Minimum similarity score (0-1)
        
        Returns:
            List of relevant facts
        """
        if not self.facts:
            return []
        
        # Encode query
        query_embedding = self.encoder.encode(query)
        
        # Calculate similarities
        similarities = []
        for i, fact_embedding in enumerate(self.embeddings):
            # Cosine similarity
            similarity = np.dot(query_embedding, fact_embedding) / (
                np.linalg.norm(query_embedding) * np.linalg.norm(fact_embedding)
            )
            similarities.append((similarity, self.facts[i]["fact"]))
        
        # Sort by similarity and filter by threshold
        similarities.sort(reverse=True, key=lambda x: x[0])
        results = [fact for score, fact in similarities[:top_k] if score >= threshold]
        
        return results

class PersonalAssistant:
    def __init__(self, api_key: str):
        self.client = Anthropic(api_key=api_key)
        self.model = "claude-sonnet-4-5"
        self.conversation = ConversationManager(max_messages=20)
        # Use semantic search instead
        self.knowledge = SemanticKnowledgeStore()
    
    # ... rest of the methods stay the same

assistant = PersonalAssistant(api_key=os.getenv("ANTHROPIC_API_KEY"))

## Store facts
assistant.chat("Remember that I'm allergic to peanuts")
assistant.chat("Remember that I prefer Italian food")
assistant.chat("Remember that I live in San Francisco")

## Semantic queries that don't match keywords exactly
print(assistant.chat("What foods should I avoid?"))
## Finds: "You're allergic to peanuts"
## Output: You should avoid peanuts and any foods containing peanuts or 
## peanut oil, as you're allergic to them...

print(assistant.chat("Where am I located?"))
## Finds: "You live in San Francisco"
## Output: You're located in San Francisco.

print(assistant.chat("What cuisine do I enjoy?"))
## Finds: "You prefer Italian food"
## Output: You enjoy Italian cuisine.

def calculate(expression: str) -> dict:
    """Calculator tool for mathematical operations."""
    try:
        result = eval(expression)
        return {"success": True, "result": result}
    except Exception as e:
        return {"success": False, "error": str(e)}

calculator_tool = {
    "name": "calculate",
    "description": "Perform mathematical calculations. Input should be a valid Python expression.",
    "input_schema": {
        "type": "object",
        "properties": {
            "expression": {
                "type": "string",
                "description": "Mathematical expression to evaluate (e.g., '2 + 2', '15 * 3.5')"
            }
        },
        "required": ["expression"]
    }
}

class PersonalAssistantWithTools:
    """Assistant with memory and tool use."""
    
    def __init__(self, api_key: str):
        self.client = Anthropic(api_key=api_key)
        self.model = "claude-sonnet-4-5"
        self.conversation = ConversationManager(max_messages=20)
        self.knowledge = SemanticKnowledgeStore()
        self.tools = {"calculate": calculate}
    
    def chat(self, user_message: str) -> str:
        """Process message with tool use support."""
        # Handle memory commands
        if self._is_memory_command(user_message):
            return self._handle_memory_command(user_message)
        
        # Search for relevant facts
        relevant_facts = self.knowledge.search(user_message)
        system_prompt = self._build_system_prompt(relevant_facts)
        
        # Add user message
        self.conversation.add_user_message(user_message)
        
        # Get response with tools
        response = self.client.messages.create(
            model=self.model,
            max_tokens=1024,
            system=system_prompt,
            tools=[calculator_tool],
            messages=self.conversation.get_messages()
        )
        
        # Handle tool use if needed
        while response.stop_reason == "tool_use":
            # Extract tool call
            tool_use_block = next(
                block for block in response.content 
                if block.type == "tool_use"
            )
            
            # Execute tool
            tool_result = self.tools[tool_use_block.name](
                **tool_use_block.input
            )
            
            # Add tool use to conversation
            self.conversation.add_assistant_message(response.content)
            
            # Add tool result
            self.conversation.messages.append({
                "role": "user",
                "content": [{
                    "type": "tool_result",
                    "tool_use_id": tool_use_block.id,
                    "content": str(tool_result)
                }]
            })
            
            # Get next response
            response = self.client.messages.create(
                model=self.model,
                max_tokens=1024,
                system=system_prompt,
                tools=[calculator_tool],
                messages=self.conversation.get_messages()
            )
        
        # Extract final response
        assistant_message = response.content[0].text
        self.conversation.add_assistant_message(assistant_message)
        
        return assistant_message
    
    # ... other methods same as before

assistant = PersonalAssistantWithTools(api_key=os.getenv("ANTHROPIC_API_KEY"))

## Store a fact
print(assistant.chat("Remember that I need to save \$500 per month"))
## Output: Got it! I'll remember that you need to save $500 per month

## Use tools with memory
print(assistant.chat("I earn \$3000 per month. After my savings goal, how much do I have left?"))
## Agent retrieves: "You need to save $500 per month"
## Agent uses calculator: 3000 - 500
## Output: After setting aside your $500 monthly savings goal, you'll have 
## $2,500 left for other expenses.

print(assistant.chat("If I split the remaining amount across 4 weeks, how much is that per week?"))
## Agent remembers previous calculation: $2,500
## Agent uses calculator: 2500 / 4
## Output: Splitting your remaining $2,500 across 4 weeks gives you $625 per week.

class PersonalAssistantWithManagement(PersonalAssistantWithTools):
    """Assistant with memory management capabilities."""
    
    def list_facts(self) -> str:
        """Show all stored facts."""
        facts = self.knowledge.get_all_facts()
        
        if not facts:
            return "I don't have any facts stored yet."
        
        facts_list = "\n".join(f"{i+1}. {fact}" for i, fact in enumerate(facts))
        return f"Here's what I know about you:\n\n{facts_list}"
    
    def forget_fact(self, fact_number: int) -> str:
        """Remove a specific fact."""
        facts = self.knowledge.facts
        
        if fact_number < 1 or fact_number > len(facts):
            return f"Invalid fact number. I have {len(facts)} facts stored."
        
        removed_fact = facts[fact_number - 1]["fact"]
        
        # Remove from both lists
        facts.pop(fact_number - 1)
        self.knowledge.embeddings.pop(fact_number - 1)
        self.knowledge._save_data()
        
        return f"Forgot: {removed_fact}"
    
    def update_fact(self, fact_number: int, new_fact: str) -> str:
        """Update an existing fact."""
        facts = self.knowledge.facts
        
        if fact_number < 1 or fact_number > len(facts):
            return f"Invalid fact number. I have {len(facts)} facts stored."
        
        old_fact = facts[fact_number - 1]["fact"]
        
        # Update fact and embedding
        facts[fact_number - 1]["fact"] = new_fact
        facts[fact_number - 1]["timestamp"] = datetime.now().isoformat()
        self.knowledge.embeddings[fact_number - 1] = self.knowledge.encoder.encode(new_fact)
        self.knowledge._save_data()
        
        return f"Updated: '{old_fact}' → '{new_fact}'"

assistant = PersonalAssistantWithManagement(api_key=os.getenv("ANTHROPIC_API_KEY"))

## Store some facts
assistant.chat("Remember that I live in San Francisco")
assistant.chat("Remember that I'm allergic to peanuts")
assistant.chat("Remember that I prefer Italian food")

## List all facts
print(assistant.list_facts())
## Output:
## Here's what I know about you:
## 
## 1. You live in San Francisco
## 2. You're allergic to peanuts
## 3. You prefer Italian food

## Update a fact
print(assistant.update_fact(1, "You live in Oakland"))
## Output: Updated: 'You live in San Francisco' → 'You live in Oakland'

## Remove a fact
print(assistant.forget_fact(2))
## Output: Forgot: You're allergic to peanuts

def test_memory_system():
    """Test that memory works correctly."""
    assistant = PersonalAssistant(api_key=os.getenv("ANTHROPIC_API_KEY"))
    
    # Test 1: Store and retrieve
    assistant.chat("Remember that my birthday is July 20")
    response = assistant.chat("When is my birthday?")
    assert "july 20" in response.lower(), "Failed to retrieve birthday"
    
    # Test 2: Semantic search
    assistant.chat("Remember that I'm allergic to peanuts")
    response = assistant.chat("What should I avoid eating?")
    assert "peanut" in response.lower(), "Failed semantic search"
    
    # Test 3: Conversation context
    assistant.chat("I'm planning a trip to Japan")
    response = assistant.chat("What's the best time to visit?")
    assert "japan" in response.lower(), "Lost conversation context"
    
    # Test 4: Persistence
    assistant2 = PersonalAssistant(api_key=os.getenv("ANTHROPIC_API_KEY"))
    response = assistant2.chat("When is my birthday?")
    assert "july 20" in response.lower(), "Facts not persisted"
    
    print("All tests passed!")

test_memory_system()