def initialize_pyrit()

in pyrit/common/initialization.py [0:0]

• 23 lines of code
• 4 McCabe index (conditional complexity)

def initialize_pyrit(memory_db_type: Union[MemoryDatabaseType, str], **memory_instance_kwargs: Optional[Any]) -> None:
    """
    Initializes PyRIT with the provided memory instance and loads environment files.

    Args:
        memory_db_type (MemoryDatabaseType): The MemoryDatabaseType string literal which indicates the memory
            instance to use for central memory. Options include "InMemory", "DuckDB", and "AzureSQL".
        **memory_instance_kwargs (Optional[Any]): Additional keyword arguments to pass to the memory instance.
    """
    if memory_db_type not in get_args(MemoryDatabaseType):
        raise ValueError(
            f"Memory database type '{memory_db_type}' is not a supported type {get_args(MemoryDatabaseType)}"
        )

    _load_environment_files()

    from pyrit.memory import (
        AzureSQLMemory,
        CentralMemory,
        DuckDBMemory,
        MemoryInterface,
    )

    memory: MemoryInterface = None
    if memory_db_type == IN_MEMORY:
        logger.info("Using in-memory DuckDB database.")
        memory = DuckDBMemory(db_path=":memory:", **memory_instance_kwargs)
    elif memory_db_type == DUCK_DB:
        logger.info("Using persistent DuckDB database.")
        memory = DuckDBMemory(**memory_instance_kwargs)
    else:
        logger.info("Using AzureSQL database.")
        memory = AzureSQLMemory(**memory_instance_kwargs)

    CentralMemory.set_memory_instance(memory)