Implement Episodic Memory with Retrieval

Medium
Agents

Episodic Memory in AI Agents

Episodic memory stores past experiences (task + steps + outcome) for future reference.

Task

Implement EpisodicMemory that:

  1. Stores episodes with embeddings for semantic retrieval.
  2. Retrieves most similar past episodes for a new task.
  3. Extracts 'lessons' from similar episodes (especially failures).
  4. Evicts oldest episodes when capacity exceeded.

Constraints

  • Eviction strategy: LRU (evict oldest by timestamp).
  • Lessons from failed episodes: 'Avoid: {task}' — from successful ones: 'Replicate: {task}'.
  • Similarity computed via cosine similarity on embeddings.

Examples

Example 1:
Input: mem.store(Episode('1', 'debug async', [...], 'success', ..., ['async'])) mem.get_lessons('fix async bug')
Output: ['Replicate: debug async']
Explanation: Similar past episode was successful, so lesson is to replicate approach.

Starter Code

from typing import List, Dict, Any, Optional
from datetime import datetime
import json

@dataclass
class Episode:
    episode_id: str
    task: str
    steps: List[Dict]
    outcome: str  # 'success' | 'failure' | 'partial'
    timestamp: str
    tags: List[str]
    embedding: Optional[List[float]] = None

class EpisodicMemory:
    def __init__(self, embed_fn: callable, max_episodes: int = 100):
        self.embed_fn = embed_fn
        self.max_episodes = max_episodes
        self.episodes: List[Episode] = []

    def store(self, episode: Episode) -> None:
        # TODO: Store with eviction if over limit
        pass

    def retrieve_similar(self, task: str, k: int = 3) -> List[Episode]:
        # TODO: Find most similar past episodes
        pass

    def get_lessons(self, task: str) -> List[str]:
        # TODO: Extract actionable lessons from similar episodes
        pass

    def evict_oldest(self, n: int = 1) -> None:
        pass
Lines: 1Characters: 0
Ready
The AI Interview - Master AI/ML Interviews