6. Azure SQL Memory#

The memory AzureSQL database can be thought of as a normalized source of truth. The memory module is the primary way pyrit keeps track of requests and responses to targets and scores. Most of this is done automatically. All orchestrators write to memory for later retrieval. All scorers also write to memory when scoring.

The schema is found in memory_models.py and can be programatically viewed as follows

Azure Login#

PyRIT AzureSQLMemory supports only Azure Entra ID authentication at this time. User ID/password-based login is not available.

Please log in to your Azure account before running this notebook:

  • Log in with the proper scope to obtain the correct access token:

    az login --scope https://database.windows.net//.default
    

Environment Variables#

Please set the following environment variables to run AzureSQLMemory interactions:

  • AZURE_SQL_DB_CONNECTION_STRING = “

  • AZURE_STORAGE_ACCOUNT_RESULTS_CONTAINER_URL = “” (which uses delegation SAS) but needs login to Azure.

To use regular key-based authentication, please also set:

  • AZURE_STORAGE_ACCOUNT_RESULTS_SAS_TOKEN

from pyrit.memory import AzureSQLMemory
from pyrit.common import default_values

default_values.load_environment_files()
memory = AzureSQLMemory()

memory.print_schema()
Schema for EmbeddingData:
  Column id (UNIQUEIDENTIFIER)
  Column embedding (NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column embedding_type_name (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
Schema for PromptMemoryEntries:
  Column id (UNIQUEIDENTIFIER)
  Column role (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column conversation_id (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column sequence (INTEGER)
  Column timestamp (DATETIME)
  Column labels (NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column prompt_metadata (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column converter_identifiers (NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column prompt_target_identifier (NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column orchestrator_identifier (NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column response_error (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column original_value_data_type (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column original_value (NVARCHAR(4000) COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column original_value_sha256 (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column converted_value_data_type (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column converted_value (NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column converted_value_sha256 (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column original_prompt_id (UNIQUEIDENTIFIER)
Schema for ScoreEntries:
  Column id (UNIQUEIDENTIFIER)
  Column score_value (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column score_value_description (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column score_type (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column score_category (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column score_rationale (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column score_metadata (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column scorer_class_identifier (NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column prompt_request_response_id (UNIQUEIDENTIFIER)
  Column timestamp (DATETIME)
  Column task (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
Schema for SeedPromptEntries:
  Column id (UNIQUEIDENTIFIER)
  Column value (NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column data_type (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column name (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column dataset_name (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column harm_categories (NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column description (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column authors (NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column groups (NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column source (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column date_added (DATETIME)
  Column added_by (VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column prompt_metadata (NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column parameters (NVARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS")
  Column prompt_group_id (UNIQUEIDENTIFIER)
  Column sequence (INTEGER)

Basic Azure SQL Memory Programming Usage#

The pyrit.memory.azure_sql_memory module provides functionality to keep track of the conversation history, scoring, data, and more using Azure SQL. You can use memory to read and write data. Here is an example that retrieves a normalized conversation:

from uuid import uuid4
from pyrit.memory import AzureSQLMemory
from pyrit.models import PromptRequestPiece, PromptRequestResponse


conversation_id = str(uuid4())

message_list = [
    PromptRequestPiece(
        role="user", original_value="Hi, chat bot! This is my initial prompt.", conversation_id=conversation_id
    ),
    PromptRequestPiece(
        role="assistant", original_value="Nice to meet you! This is my response.", conversation_id=conversation_id
    ),
    PromptRequestPiece(
        role="user",
        original_value="Wonderful! This is my second prompt to the chat bot!",
        conversation_id=conversation_id,
    ),
]

memory = AzureSQLMemory()

memory.add_request_response_to_memory(request=PromptRequestResponse([message_list[0]]))
memory.add_request_response_to_memory(request=PromptRequestResponse([message_list[1]]))
memory.add_request_response_to_memory(request=PromptRequestResponse([message_list[2]]))


entries = memory.get_conversation(conversation_id=conversation_id)

for entry in entries:
    print(entry)


# Cleanup memory resources
memory.dispose_engine()
None: user: Hi, chat bot! This is my initial prompt.
None: assistant: Nice to meet you! This is my response.
None: user: Wonderful! This is my second prompt to the chat bot!