Refactor vector storage function and enhance logging in VectorStore

This commit refines the `store_vector` function within the VectorStore class by improving its structure and adding detailed docstrings for better clarity. It also enhances logging for vector storage operations, providing more informative debug messages for different tiers (STM, IM, LTM). Additionally, the logging utility is updated to allow for customizable log directory paths, ensuring more flexible logging configurations.

Author: csmangum

memory/embeddings/vector_store.py

@@ -556,26 +556,58 @@ def store_memory_vectors(
         embeddings = memory_entry.get("embeddings", {})
         metadata = memory_entry.get("metadata", {})
 
-        def store_vector(index, memory_id: str, vector: List[float], metadata: Dict[str, Any], tier: str) -> bool:
+        def store_vector(
+            index,
+            memory_id: str,
+            vector: List[float],
+            metadata: Dict[str, Any],
+            tier: str,
+        ) -> bool:
+            """
+            Store a vector in the appropriate index.
+
+            Args:
+                index: Vector index to store in
+                memory_id: Unique identifier for the memory
+                vector: Vector to store
+                metadata: Metadata to store
+                tier: Tier to store the vector in ("stm", "im", or "ltm")
+
+            Returns:
+                True if storage was successful
+            """
             try:
                 logger.debug(f"Storing {tier.upper()} vector for memory {memory_id}")
                 return index.add(memory_id, vector, metadata)
             except Exception as e:
-                logger.error(f"Failed to store {tier.upper()} vector for memory {memory_id}: {e}")
+                logger.error(
+                    f"Failed to store {tier.upper()} vector for memory {memory_id}: {e}"
+                )
                 return False
 
         success = True
 
         # Store in appropriate index based on tier
         if tier == "stm":
-            success = store_vector(self.stm_index, memory_id, embeddings["full_vector"], metadata, "stm")
+            success = store_vector(
+                self.stm_index, memory_id, embeddings["full_vector"], metadata, "stm"
+            )
         elif tier == "im":
             #! TODO: Use compressed vector
-            success = store_vector(self.im_index, memory_id, embeddings["full_vector"], metadata, "im")
-            logger.debug(f"Result of storing IM vector: {success}")
+            logger.debug(
+                f"@@@@@@@@@@@@@@@@@@@@@@@@@ Storing IM vector for memory {memory_id}"
+            )
+            success = store_vector(
+                self.im_index, memory_id, embeddings["full_vector"], metadata, "im"
+            )
+            logger.debug(
+                f"@@@@@@@@@@@@@@@@@@@@@@@@@ Result of storing IM vector: {success}"
+            )
         elif tier == "ltm":
             #! TODO: Use abstract vector
-            success = store_vector(self.ltm_index, memory_id, embeddings["full_vector"], metadata, "ltm")
+            success = store_vector(
+                self.ltm_index, memory_id, embeddings["full_vector"], metadata, "ltm"
+            )
 
         return success
 

memory/utils/logging.py

@@ -8,8 +8,11 @@
 print(f"Current working directory: {os.getcwd()}")
 print(f"Logs directory: {LOGS_DIR}")
 
+# Configure logging
+logger = logging.getLogger()
 
-def setup_logging(demo_name: str) -> logging.Logger:
+
+def setup_logging(demo_name: str, log_dir: str = LOGS_DIR) -> logging.Logger:
     """Set up logging to both console and file.
 
     Args:
@@ -20,11 +23,11 @@ def setup_logging(demo_name: str) -> logging.Logger:
     """
     try:
         # Create logs directory if it doesn't exist
-        os.makedirs(LOGS_DIR, exist_ok=True)
-        print(f"Created logs directory at: {LOGS_DIR}")
+        os.makedirs(log_dir, exist_ok=True)
+        print(f"Created logs directory at: {log_dir}")
 
         # Use a fixed log filename based on demo name (without timestamp)
-        log_file = os.path.join(LOGS_DIR, f"{demo_name}.log")
+        log_file = os.path.join(log_dir, f"{demo_name}.log")
         print(f"Log file will be created at: {log_file}")
 
         # Clear the existing log file if it exists
@@ -33,8 +36,7 @@ def setup_logging(demo_name: str) -> logging.Logger:
             pass
         print(f"Initialized log file: {log_file}")
 
-        # Configure logging
-        logger = logging.getLogger()
+        
         logger.setLevel(logging.DEBUG)
 
         # Clear existing handlers if any
@@ -80,4 +82,7 @@ def filter(self, record):
         print(f"Error setting up logging: {e}")
         # Fallback to basic console logging
         logging.basicConfig(level=logging.INFO)
-        return logging.getLogger()
+        return logging.getLogger()
+    
+    
+