src/utils.py

import pickle
import os

# Define the persistent storage path, navigating two levels up from src/
EMBEDDINGS_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'data', 'embeddings.pkl')

def load_embeddings():
    """
    Loads all registered embeddings and associated names from the pickle file.
    
    Returns:
        dict: A dictionary with 'embeddings' (list of numpy arrays) and 'names' (list of strings).
    """
    if not os.path.exists(EMBEDDINGS_PATH):
        # Initialize an empty structure if the file doesn't exist
        return {'embeddings': [], 'names': []}
    
    try:
        with open(EMBEDDINGS_PATH, 'rb') as f:
            data = pickle.load(f)
            # Ensure the loaded data has the expected structure
            if 'embeddings' not in data or 'names' not in data:
                 return {'embeddings': [], 'names': []}
            return data
    except Exception as e:
        print(f"Error loading embeddings: {e}")
        # In case of corruption, return empty data
        return {'embeddings': [], 'names': []}

def save_embeddings(data):
    """
    Saves the updated embeddings and names back to the pickle file.
    
    Args:
        data (dict): The dictionary containing 'embeddings' and 'names' lists.
    """
    # Ensure the data directory exists
    os.makedirs(os.path.dirname(EMBEDDINGS_PATH), exist_ok=True)
    
    try:
        with open(EMBEDDINGS_PATH, 'wb') as f:
            pickle.dump(data, f)
        return True
    except Exception as e:
        print(f"Error saving embeddings: {e}")
        return False