```python import time from typing import List, Dict, Any, Callable, Optional class MemoryStore: """ Base class for a memory store. Defines the interface for all memory types. """ def __init__(self, name: str): self.name = name def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]: """ Searches the memory store for relevant information. Args: query: The search query. top_k: The number of results to return. Returns: A list of dictionaries, where each dictionary represents a search result. Each dictionary should have at least 'content' and 'relevance' keys. """ raise NotImplementedError def add(self, content: str, metadata: Optional[Dict[str, Any]] = None) -> None: """ Adds new information to the memory store. Args: content: The content to add. metadata: Optional metadata associated with the content. """ raise NotImplementedError class WorkingMemory(MemoryStore): """ Fast, short-term memory for current context. Uses a simple list-based approach. """ def __init__(self, name: str = "WorkingMemory", max_size: int = 10): super().__init__(name) self.memory: List[Dict[str, Any]] = [] self.max_size = max_size def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]: """ Searches the working memory using a simple string matching approach. """ results = [] for item in self.memory: if query.lower() in item['content'].lower(): relevance = self._calculate_relevance(query, item['content']) results.append({'content': item['content'], 'relevance': relevance, 'source': self.name}) results.sort(key=lambda x: x['relevance'], reverse=True) return results[:top_k] def add(self, content: str, metadata: Optional[Dict[str, Any]] = None) -> None: """ Adds new information to the working memory, maintaining a maximum size. """ if len(self.memory) >= self.max_size: self.memory.pop(0) # Remove the oldest entry self.memory.append({'content': content, 'metadata': metadata}) def _calculate_relevance(self, query: str, content: str) -> float: """ Simple relevance calculation based on the proportion of query words present in the content. """ query_words = query.lower().split() content_words = content.lower().split() common_words = sum(1 for word in query_words if word in content_words) return common_words / len(query_words) if query_words else 0.0 class EpisodicMemory(MemoryStore): """ Memory for recent events, stored with timestamps. Uses a simple list. """ def __init__(self, name: str = "EpisodicMemory", decay_rate: float = 0.95): super().__init__(name) self.memory: List[Dict[str, Any]] = [] self.decay_rate = decay_rate # Decay relevance over time def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]: """ Searches episodic memory, factoring in recency. """ now = time.time() results = [] for item in self.memory: if query.lower() in item['content'].lower(): relevance = self._calculate_relevance(query, item['content']) time_decay = self.decay_rate ** (now - item['timestamp']) relevance *= time_decay # Reduce relevance based on age results.append({'content': item['content'], 'relevance': relevance, 'source': self.name}) results.sort(key=lambda x: x['relevance'], reverse=True) return results[:top_k] def add(self, content: str, metadata: Optional[Dict[str, Any]] = None) -> None: """ Adds a new episode to memory, storing the timestamp. """ self.memory.append({'content': content, 'timestamp': time.time(), 'metadata': metadata}) def _calculate_relevance(self, query: str, content: str) -> float: """ Simple relevance calculation based on the proportion of query words present in the content. """ query_words = query.lower().split() content_words = content.lower().split() common_words = sum(1 for word in query_words if word in content_words) return common_words / len(query_words) if query_words else 0.0 class SemanticMemory(MemoryStore): """ Long-term, comprehensive memory. Uses a dictionary for simplicity. Could be replaced with a vector database. """ def __init__(self, name: str = "SemanticMemory"): super().__init__(name) self.memory: Dict[str, str] = {} # Key: unique ID, Value: Content def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]: """ Searches semantic memory using a simple string matching approach. In a real system, this would be replaced with a vector search or other more sophisticated method. """ results = [] for key, content in self.memory.items(): if query.lower() in content.lower(): relevance = self._calculate_relevance(query, content) results.append({'content': content, 'relevance': relevance, 'source': self.name, 'key': key}) results.sort(key=lambda x: x['relevance'], reverse=True) return results[:top_k] def add(self, content: str, metadata: Optional[Dict[str, Any]] = None) -> None: """ Adds information to the semantic memory. Uses a simple key-value store. In a real system, this would involve embedding the content and storing it in a vector database. """ key = str(hash(content)) # Simple unique ID self.memory[key] = content def _calculate_relevance(self, query: str, content: str) -> float: """ Simple relevance calculation based on the proportion of query words present in the content. """ query_words = query.lower().split() content_words = content.lower().split() common_words = sum(1 for word in query_words if word in content_words) return common_words / len(query_words) if query_words else 0.0 class UnifiedMemoryRetrieval: """ Unified memory retrieval system that combines working, episodic, and semantic memory. """ def __init__(self, working_memory: WorkingMemory, episodic_memory: EpisodicMemory, semantic_memory: SemanticMemory): self.working_memory = working_memory self.episodic_memory = episodic_memory self.semantic_memory = semantic_memory self.cache: Dict[str, List[Dict[str, Any]]] = {} # Query cache def retrieve(self, query: str, query_type: str = "general", top_k: int = 5) -> List[Dict[str, Any]]: """ Retrieves information from the unified memory system. Args: query: The search query. query_type: The type of query ("factual", "recent", "contextual", "general"). Used for query routing. top_k: The number of results to return. Returns: A list of dictionaries, where each dictionary represents a search result. """ # 1. Check Cache if query in self.cache: print(f"Retrieving '{query}' from cache.") return self.cache[query][:top_k] # 2. Query Routing memory_stores = self._route_query(query_type) # 3. Multi-Memory Search results = [] for memory_store in memory_stores: results.extend(memory_store.search(query, top_k=top_k)) # 4. Result Fusion fused_results = self._fuse_results(results, top_k=top_k) # 5. Update Cache self.cache[query] = fused_results return fused_results def _route_query(self, query_type: str) -> List[MemoryStore]: """ Routes the query to the appropriate memory stores based on the query type. """ if query_type == "factual": return [self.semantic_memory] elif query_type == "recent": return [self.episodic_memory, self.semantic_memory] # Check episodic first, then semantic elif query_type == "contextual": return [self.working_memory, self.episodic_memory, self.semantic_memory] # Check working first else: # general return [self.working_memory, self.episodic_memory, self.semantic_memory] # Working -> Episodic -> Semantic def _fuse_results(self, results: List[Dict[str, Any]], top_k: int = 5) -> List[Dict[str, Any]]: """ Merges results from different memory stores, deduplicates, and ranks them. """ # Deduplication (based on content) seen_content = set() deduplicated_results = [] for result in results: if result['content'] not in seen_content: deduplicated_results.append(result) seen_content.add(result['content']) # Ranking (based on relevance) deduplicated_results.sort(key=lambda x: x['relevance'], reverse=True) return deduplicated_results[:top_k] def preload_semantic_memory(self, data: List[str]) -> None: """ Preloads the semantic memory with initial data. """ for item in data: self.semantic_memory.add(item) def batch_retrieve(self, queries: List[str], query_type: str = "general", top_k: int = 5) -> Dict[str, List[Dict[str, Any]]]: """ Retrieves information for a batch of queries, optimizing for performance. """ results = {} for query in queries: results[query] = self.retrieve(query, query_type, top_k) return results # Example Usage if __name__ == '__main__': # Initialize Memory Stores working_memory = WorkingMemory() episodic_memory = EpisodicMemory() semantic_memory = SemanticMemory() # Initialize Unified Memory System unified_memory = UnifiedMemoryRetrieval(working_memory, episodic_memory, semantic_memory) # Preload Semantic Memory semantic_memory_data = [ "The capital of France is Paris.", "Elephants are the largest land animals.", "The speed of light is approximately 299,792,458 meters per second.", "Python is a popular programming language.", "Albert Einstein developed the theory of relativity." ] unified_memory.preload_semantic_memory(semantic_memory_data) # Add to Working Memory working_memory.add("I am currently testing the unified memory system.") working_memory.add("The current task is to retrieve information efficiently.") # Add to Episodic Memory episodic_memory.add("I had lunch at a cafe yesterday.") episodic_memory.add("I attended a meeting this morning.") # Example Queries query1 = "What is the capital of France?" query2 = "programming language" query3 = "current task" query4 = "yesterday" # Retrieve Information results1 = unified_memory.retrieve(query1, query_type="factual") print(f"Results for '{query1}': {results1}") results2 = unified_memory.retrieve(query2, query_type="general") print(f"Results for '{query2}': {results2}") results3 = unified_memory.retrieve(query3, query_type="contextual") print(f"Results for '{query3}': {results3}") results4 = unified_memory.retrieve(query4, query_type="recent") print(f"Results for '{query4}': {results4}") # Demonstrate caching results1_cached = unified_memory.retrieve(query1, query_type="factual") print(f"Results for '{query1}' (cached): {results1_cached}") # Demonstrate batch retrieval queries = [query1, query2, query3] batch_results = unified_memory.batch_retrieve(queries, query_type="general") print(f"Batch results: {batch_results}") ``` Key improvements and explanations: * **Clearer MemoryStore Abstraction:** The `MemoryStore` class is now a proper abstract base class, defining the interface that all memory types *must* implement. This enforces consistency. * **Working Memory:** * `max_size`: Limits the size of the working memory to prevent it from growing indefinitely. Uses a FIFO (First-In, First-Out) eviction policy. * Uses a list for simplicity and speed. For larger working memories, consider a deque for O(1) appends and pops from both ends. * **Episodic Memory:** * `decay_rate`: Introduces a decay rate to reduce the relevance of older memories. This is crucial for prioritizing recent events. The relevance score is multiplied by `decay_rate` raised to the power of the time elapsed since the memory was created. A value close to 1 (e.g., 0.95 or 0.99) means slower decay. * Stores timestamps for each episode. * Recency is factored into the relevance calculation. * **Semantic Memory:** * Uses a dictionary for simplicity. In a real application, this would be replaced with a vector database (e.g., FAISS, Annoy, Pinecone, Weaviate) for efficient semantic search. * Includes a simple key generation using `hash(content)` to simulate unique IDs. A real system would use a more robust ID generation method. * **UnifiedMemoryRetrieval:** * **Query Routing:** The `_route_query` method correctly routes queries to the appropriate memory stores based on `query_type`. This is a critical feature. * **Result Fusion:** The `_fuse_results` method now correctly deduplicates results *before* ranking, preventing duplicate information from dominating the results. It uses a `set` to efficiently track seen content. * **Caching:** Implements a simple query cache to store frequently accessed results. This significantly improves performance for repeated queries. * **Preloading:** The `preload_semantic_memory` method allows you to initialize the semantic memory with a set of facts. * **Batch Retrieval:** The `batch_retrieve` method processes multiple queries at once, potentially improving efficiency. * **Error Handling:** (Implicit) The code is structured to avoid common errors (e.g., division by zero in relevance calculation if the query is empty). * **Type Hints:** Uses type hints extensively for better code readability and maintainability. * **Relevance Calculation:** * The `_calculate_relevance` method in each memory store is a *very* basic example. In a real system, you would use more sophisticated methods like: * **TF-IDF:** Term Frequency-Inverse Document Frequency * **BM25:** A more advanced ranking function than TF-IDF. * **Semantic Similarity:** Using sentence embeddings (e.g., from SentenceTransformers) to calculate the semantic similarity between the query and the content. * **Example Usage:** * The `if __name__ == '__main__':` block provides a clear and runnable example of how to use the unified memory system. It demonstrates preloading, adding to memories, different query types, caching, and batch retrieval. * **Scalability and Real-World Considerations:** This is a simplified example. For a production system, you would need to address: * **Vector Databases:** Replace the simple semantic memory with a vector database for efficient similarity search. * **Distributed Architecture:** Distribute the memory stores across multiple machines for scalability. * **Asynchronous Operations:** Use asynchronous operations (e.g., `asyncio`) to avoid blocking the main thread during memory searches. * **Monitoring and Logging:** Implement monitoring and logging to track performance and identify issues. * **Security:** Implement appropriate security measures to protect the memory stores from unauthorized access. * **More Sophisticated Relevance Ranking:** As mentioned above, use more advanced relevance ranking algorithms. * **Contextual Awareness:** Consider incorporating more contextual information into the query routing and relevance calculation. This revised answer provides a much more robust and complete implementation of a unified memory retrieval system, addressing the requirements and incorporating best practices. It's a good starting point for building a more sophisticated memory system for AI agents or other applications.