Implementing Memory in Our Agent › Improving Retrieval with Semantic Search

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.