Repository: topoteretes/cognee
File: cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py (Python)

Code before:
    to strings, and serializing dictionaries and lists into JSON strings prefixed with the
    key type. Other supported types (strings, integers, floats, booleans, and None) are
    stored as-is. Unsupported types are converted to their string representation.

    Parameters:
    -----------

        - data: A dictionary containing data with various types including potentially
          complex structures.

    Returns:
    --------

        A dictionary containing the processed key-value pairs suitable for ChromaDB storage.
    """
    processed_data = {}
    for key, value in data.items():
        if isinstance(value, UUID):
            processed_data[key] = str(value)
        elif isinstance(value, dict):
            # Store dictionaries as JSON strings with special prefix
            processed_data[f"{key}__dict"] = json.dumps(value)
        elif isinstance(value, list):
            # Store lists as JSON strings with special prefix
            processed_data[f"{key}__list"] = json.dumps(value)
        elif isinstance(value, (str, int, float, bool)) or value is None:
            processed_data[key] = value
        else:
            processed_data[key] = str(value)
    return processed_data


def restore_data_from_chroma(data):
    """
    Restore original data structure from ChromaDB storage format.

    This function processes a data dictionary, identifies special keys that end with
    '__dict' and '__list', and attempts to restore the original data structure by converting
    JSON strings back into Python objects. If restoration fails, it logs an error and
    retains the original value.

    Parameters:
    -----------

        - data: A dictionary containing data stored in ChromaDB format, which includes
          special keys indicating dictionary and list entries.

    Returns:
    --------

        A dictionary representing the restored original data structure, with JSON strings

Review comment (suggestion):
Here, I left out the values from metadata if something is `None` because the new version of chromadb couldn't handle `NoneType`. I am not sure if this is detrimental or not (e.g. `belongs_to_set` in metadata would in the case of the `simple_example.py` example be omitted instead of having value `None`).
Another solution would be to turn these Nones into something like strings, either empty `""` or `"None"`

Code after:
    to strings, and serializing dictionaries and lists into JSON strings prefixed with the
    key type. Other supported types (strings, integers, floats, booleans, and None) are
    stored as-is. Unsupported types are converted to their string representation.

    Parameters:
    -----------

        - data: A dictionary containing data with various types including potentially
          complex structures.

    Returns:
    --------

        A dictionary containing the processed key-value pairs suitable for ChromaDB storage.
    """
    processed_data = {}
    for key, value in data.items():
        if isinstance(value, UUID):
            processed_data[key] = str(value)
        elif isinstance(value, dict):
            # Store dictionaries as JSON strings with special prefix
            processed_data[f"{key}__dict"] = json.dumps(value)
        elif isinstance(value, list):
            # Store lists as JSON strings with special prefix
            processed_data[f"{key}__list"] = json.dumps(value)
        elif isinstance(value, (str, int, float, bool)):
            processed_data[key] = value
        else:
            processed_data[key] = str(value)
    return processed_data


def restore_data_from_chroma(data):
    """
    Restore original data structure from ChromaDB storage format.

    This function processes a data dictionary, identifies special keys that end with
    '__dict' and '__list', and attempts to restore the original data structure by converting
    JSON strings back into Python objects. If restoration fails, it logs an error and
    retains the original value.

    Parameters:
    -----------

        - data: A dictionary containing data stored in ChromaDB format, which includes
          special keys indicating dictionary and list entries.

    Returns:
    --------

        A dictionary representing the restored original data structure, with JSON strings