Merge pull request #127 from Dooders/similarity-search-validation

Similarity search validation

Author: csmangum

memory/embeddings/vector_store.py

@@ -243,7 +243,9 @@ def _ensure_index(self) -> None:
             indices = self.redis.execute_command("FT._LIST")
             # Convert both the index name and list items to strings for comparison
             index_name_str = str(self.index_name)
-            indices_str = [str(idx) if isinstance(idx, bytes) else idx for idx in indices]
+            indices_str = [
+                str(idx) if isinstance(idx, bytes) else idx for idx in indices
+            ]
             if index_name_str in indices_str:
                 self._index_exists = True
                 return
@@ -534,11 +536,14 @@ def __init__(
             "Redis" if redis_client else "in-memory",
         )
 
-    def store_memory_vectors(self, memory_entry: Dict[str, Any]) -> bool:
+    def store_memory_vectors(
+        self, memory_entry: Dict[str, Any], tier: str = "stm"
+    ) -> bool:
         """Store vectors for a memory entry in appropriate indices.
 
         Args:
             memory_entry: Memory entry with embeddings
+            tier: Tier to store the vectors in ("stm", "im", or "ltm")
 
         Returns:
             True if storage was successful
@@ -553,22 +558,27 @@ def store_memory_vectors(self, memory_entry: Dict[str, Any]) -> bool:
 
         success = True
 
-        # Store in STM index if full vector is available
-        if "full_vector" in embeddings:
+        # Store in STM index
+        if tier == "stm":
+            logger.debug("Storing STM vector for memory %s", memory_id)
             success = success and self.stm_index.add(
                 memory_id, embeddings["full_vector"], metadata
             )
 
-        # Store in IM index if compressed vector is available
-        if "compressed_vector" in embeddings:
+        # Store in IM index
+        if tier == "im":
+            #! TODO: Use compressed vector
+            logger.debug("Storing IM vector for memory %s", memory_id)
             success = success and self.im_index.add(
-                memory_id, embeddings["compressed_vector"], metadata
+                memory_id, embeddings["full_vector"], metadata
             )
 
-        # Store in LTM index if abstract vector is available
-        if "abstract_vector" in embeddings:
+        # Store in LTM index
+        if tier == "ltm":
+            #! TODO: Use abstract vector
+            logger.debug("Storing LTM vector for memory %s", memory_id)
             success = success and self.ltm_index.add(
-                memory_id, embeddings["abstract_vector"], metadata
+                memory_id, embeddings["full_vector"], metadata
             )
 
         return success
@@ -594,36 +604,55 @@ def find_similar_memories(
         # Create filter function if metadata filter is provided
         filter_fn = None
         if metadata_filter:
-            logger.debug("Creating filter function for metadata filter: %s", metadata_filter)
+            logger.debug(
+                "Creating filter function for metadata filter: %s", metadata_filter
+            )
 
             def filter_fn(metadata):
                 logger.debug("Checking metadata: %s", metadata)
+                unmatched_keys = []  # Initialize the list before using it
                 for key, value in metadata_filter.items():
                     # Try direct match in top-level metadata
                     if key in metadata and metadata[key] == value:
                         logger.debug("Found direct match for %s: %s", key, value)
                         continue
-                    
+
                     # Special handling for 'type' field - also check 'memory_type'
-                    if key == 'type' and 'memory_type' in metadata and metadata['memory_type'] == value:
+                    if (
+                        key == "type"
+                        and "memory_type" in metadata
+                        and metadata["memory_type"] == value
+                    ):
                         logger.debug("Found match for type in memory_type: %s", value)
                         continue
-                    
+
                     # Try match in nested content.metadata
-                    if 'content' in metadata and isinstance(metadata['content'], dict):
-                        content = metadata['content']
-                        if 'metadata' in content and isinstance(content['metadata'], dict):
-                            content_metadata = content['metadata']
-                            if key in content_metadata and content_metadata[key] == value:
-                                logger.debug("Found nested match for %s: %s in content.metadata", key, value)
+                    if "content" in metadata and isinstance(metadata["content"], dict):
+                        content = metadata["content"]
+                        if "metadata" in content and isinstance(
+                            content["metadata"], dict
+                        ):
+                            content_metadata = content["metadata"]
+                            if (
+                                key in content_metadata
+                                and content_metadata[key] == value
+                            ):
+                                logger.debug(
+                                    "Found nested match for %s: %s in content.metadata",
+                                    key,
+                                    value,
+                                )
                                 continue
-                    
+
                     # No match found for this key
                     unmatched_keys.append((key, value))
                     return False
-                
+
                 if unmatched_keys:
-                    logger.debug("No matches found for the following keys and values: %s", unmatched_keys)
+                    logger.debug(
+                        "No matches found for the following keys and values: %s",
+                        unmatched_keys,
+                    )
                 else:
                     logger.debug("All filter criteria matched")
                 # All keys matched

memory/search/strategies/similarity.py

@@ -111,6 +111,14 @@ def search(
 
             # Generate query vector from input
             query_vector = self._generate_query_vector(query, current_tier)
+            
+            # Add detailed logging for vector generation
+            logger.debug(
+                "Query vector generation for tier %s - Input: %s, Output: %s",
+                current_tier,
+                query,
+                query_vector
+            )
 
             # Skip if vector generation failed
             if query_vector is None:
@@ -127,21 +135,30 @@ def search(
 
             # Find similar vectors
             logger.debug(
-                "Calling vector_store.find_similar_memories for tier %s", current_tier
-            )
-            similar_vectors = self.vector_store.find_similar_memories(
-                query_vector,
-                tier=current_tier,
-                limit=limit * 2,  # Get extra results to allow for score filtering
-                metadata_filter=metadata_filter or {},
-            )
-
-            logger.debug(
-                "Vector store returned %d results for tier %s. Raw results: %s",
-                len(similar_vectors),
+                "About to call vector_store.find_similar_memories for tier %s with vector: %s", 
                 current_tier,
-                similar_vectors,
+                query_vector
             )
+            try:
+                similar_vectors = self.vector_store.find_similar_memories(
+                    query_vector,
+                    tier=current_tier,
+                    limit=limit * 2,  # Get extra results to allow for score filtering
+                    metadata_filter=metadata_filter or {},
+                )
+                logger.debug(
+                    "Vector store returned %d results for tier %s. Raw results with scores: %s",
+                    len(similar_vectors),
+                    current_tier,
+                    [(v["id"], v["score"]) for v in similar_vectors],
+                )
+            except Exception as e:
+                logger.error(
+                    "Error in vector_store.find_similar_memories for tier %s: %s",
+                    current_tier,
+                    str(e)
+                )
+                continue
 
             # Filter by score
             filtered_vectors = [v for v in similar_vectors if v["score"] >= min_score]
@@ -150,7 +167,7 @@ def search(
                 min_score,
                 len(filtered_vectors),
                 current_tier,
-                filtered_vectors,
+                [(v["id"], v["score"]) for v in filtered_vectors],
             )
 
             # Limit results
@@ -274,22 +291,27 @@ def _generate_query_vector(
             logger.debug(
                 "Encoding dictionary query for tier %s. Query dict: %s", tier, query
             )
-            if tier == "stm":
-                vector = self.embedding_engine.encode_stm(query)
-            elif tier == "im":
-                vector = self.embedding_engine.encode_im(query)
-            elif tier == "ltm":
-                vector = self.embedding_engine.encode_ltm(query)
-
-            if vector is not None:
-                logger.debug(
-                    "Successfully generated vector of length %d: %s",
-                    len(vector),
-                    vector,
-                )
-            else:
-                logger.warning("Failed to generate vector for tier %s", tier)
-            return vector
+            try:
+                if tier == "stm":
+                    vector = self.embedding_engine.encode_stm(query)
+                elif tier == "im":
+                    vector = self.embedding_engine.encode_im(query)
+                elif tier == "ltm":
+                    vector = self.embedding_engine.encode_ltm(query)
+
+                if vector is not None:
+                    logger.debug(
+                        "Successfully generated vector of length %d for tier %s: %s",
+                        len(vector),
+                        tier,
+                        vector,
+                    )
+                else:
+                    logger.warning("Failed to generate vector for tier %s - encoding returned None", tier)
+                return vector
+            except Exception as e:
+                logger.error("Error generating vector for tier %s: %s", tier, str(e))
+                return None
 
         # If we get here, we couldn't generate a vector
         logger.warning("Could not generate vector for query type: %s", type(query))

memory/storage/mockredis/core.py

@@ -996,11 +996,47 @@ def execute_command(self, command, *args, **kwargs):
                 if "_vector_idx" in index_name and has_vector_query:
                     # Simulate a vector search by returning some random keys
                     agent_prefix = index_name.split('_')[0]
+                    metadata_filter = {}  # Initialize empty metadata filter
+                    
+                    # Extract metadata filter from args if present
+                    for i in range(len(args)):
+                        if args[i].lower() == "filter":
+                            try:
+                                metadata_filter = json.loads(args[i + 2])
+                            except:
+                                pass
+                            break
+                            
                     for key in self.store.keys():
                         if isinstance(key, str) and key.startswith(f"{agent_prefix}-"):
-                            matching_keys.append(key)
-                            if len(matching_keys) >= vector_k:
-                                break
+                            # Get the memory data to check metadata
+                            memory_data = self.store.get(key)
+                            if memory_data and isinstance(memory_data, dict):
+                                # Check if memory has the required metadata
+                                metadata = memory_data.get('metadata', {})
+                                content = memory_data.get('content', {})
+                                content_metadata = content.get('metadata', {}) if isinstance(content, dict) else {}
+                                
+                                # Check if memory matches metadata filter
+                                matches_filter = True
+                                for filter_key, filter_value in metadata_filter.items():
+                                    # Check in top-level metadata
+                                    if filter_key in metadata and metadata[filter_key] == filter_value:
+                                        continue
+                                    # Check in memory_type
+                                    if filter_key == 'type' and 'memory_type' in metadata and metadata['memory_type'] == filter_value:
+                                        continue
+                                    # Check in content.metadata
+                                    if filter_key in content_metadata and content_metadata[filter_key] == filter_value:
+                                        continue
+                                    # No match found
+                                    matches_filter = False
+                                    break
+                                
+                                if matches_filter:
+                                    matching_keys.append(key)
+                                    if len(matching_keys) >= vector_k:
+                                        break
 
                     # Simulate scores
                     scores = [random.uniform(0.5, 1.0) for _ in matching_keys]

memory/storage/redis_im.py

@@ -288,13 +288,14 @@ def store(
             return False
 
         # Verify compression level
-        compression_level = memory_entry.get("metadata", {}).get("compression_level")
-        if compression_level != 1:
-            logger.error(
-                "Invalid compression level for IM storage: %s. Expected level 1.",
-                compression_level,
-            )
-            return False
+        #! TODO: Use compression level validation
+        # compression_level = memory_entry.get("metadata", {}).get("compression_level")
+        # if compression_level != 1:
+        #     logger.error(
+        #         "Invalid compression level for IM storage: %s. Expected level 1.",
+        #         compression_level,
+        #     )
+        #     return False
 
         # Use store_with_retry for resilient storage
         return self.redis.store_with_retry(
@@ -485,7 +486,9 @@ def _store_memory_entry(self, agent_id: str, memory_entry: MemoryEntry) -> bool:
                         return False
                     return True
                 except (RedisUnavailableError, RedisTimeoutError) as e:
-                    logger.debug(f"Caught Redis error in pipeline block: {type(e).__name__}")
+                    logger.debug(
+                        f"Caught Redis error in pipeline block: {type(e).__name__}"
+                    )
                     raise  # propagate these errors
                 except redis.RedisError as e:
                     logger.exception(