pyrit.memory.AzureSQLMemory

class AzureSQLMemory(*args, **kwargs)

Bases: MemoryInterface

A class to manage conversation memory using Azure SQL Server as the backend database. It leverages SQLAlchemy Base models for creating tables and provides CRUD operations to interact with the tables.

This class encapsulates the setup of the database connection, table creation based on SQLAlchemy models, and session management to perform database operations.

__init__(*, connection_string: str | None = None, container_url: str | None = None, sas_token: str | None = None, verbose: bool = False)

Methods

__init__(*[, connection_string, ...])
add_request_pieces_to_memory(*, request_pieces)	Inserts a list of prompt request pieces into the memory storage.
add_request_response_to_memory(*, request)	Inserts a list of prompt request pieces into the memory storage.
add_scores_to_memory(*, scores)	Inserts a list of scores into the memory storage.
add_seed_prompt_groups_to_memory(*, ...[, ...])	Inserts a list of seed prompt groups into the memory storage.
add_seed_prompts_to_memory(*, prompts[, ...])	Inserts a list of prompts into the memory storage.
disable_embedding()
dispose_engine()	Dispose the engine and clean up resources.
duplicate_conversation(*, conversation_id[, ...])	Duplicates a conversation for reuse
duplicate_conversation_excluding_last_turn(*, ...)	Duplicate a conversation, excluding the last turn.
enable_embedding([embedding_model])
export_all_conversations_with_scores(*[, ...])	Exports all conversations with scores to a specified file.
export_conversation_by_id(*, conversation_id)	Exports conversation data with the given conversation ID to a specified file.
export_conversation_by_orchestrator_id(*, ...)	Exports conversation data with the given orchestrator ID to a specified file.
get_all_embeddings()	Fetches all entries from the specified table and returns them as model instances.
get_all_prompt_pieces()	Fetches all entries from the specified table and returns them as model instances.
get_chat_messages_with_conversation_id(*, ...)	Returns the memory for a given conversation_id.
get_conversation(*, conversation_id)	Retrieves a list of PromptRequestResponse objects that have the specified conversation ID.
get_prompt_ids_by_orchestrator(*, ...)
get_prompt_request_piece_by_memory_labels(*)	Retrieves a list of PromptRequestPiece objects that have the specified memory labels.
get_prompt_request_piece_by_orchestrator_id(*, ...)	Retrieves a list of PromptRequestPiece objects that have the specified orchestrator ID.
get_prompt_request_pieces_by_id(*, prompt_ids)	Retrieves a list of PromptRequestPiece objects that have the specified prompt ids.
get_scores_by_memory_labels(*, memory_labels)	Retrieves a list of Score objects associated with the PromptRequestPiece objects which have the specified memory labels.
get_scores_by_orchestrator_id(*, orchestrator_id)	Retrieves a list of Score objects associated with the PromptRequestPiece objects which have the specified orchestrator ID.
get_scores_by_prompt_ids(*, ...)	Gets a list of scores based on prompt_request_response_ids.
get_seed_prompt_dataset_names()	Returns a list of all seed prompt dataset names in the memory storage.
get_seed_prompt_groups(*[, dataset_name, ...])	Retrieves groups of seed prompts based on the provided filtering criteria._summary_
get_seed_prompts(*[, value, dataset_name, ...])	Retrieves a list of seed prompts based on the specified filters.
get_session()	Provides a session for database operations.
insert_entries(*, entries)	Inserts multiple entries into the database.
insert_entry(entry)	Inserts an entry into the Table.
print_schema()	Prints the schema of all tables in the Azure SQL database.
query_entries(model, *[, conditions, distinct])	Fetches data from the specified table model with optional conditions.
reset_database()	Drop and recreate existing tables
update_entries(*, entries, update_fields)	Updates the given entries with the specified field values.
update_labels_by_conversation_id(*, ...)	Updates the labels of prompt entries in memory for a given conversation ID.
update_prompt_entries_by_conversation_id(*, ...)	Updates prompt entries for a given conversation ID with the specified field values.
update_prompt_metadata_by_conversation_id(*, ...)	Updates the metadata of prompt entries in memory for a given conversation ID.

Attributes

AZURE_SQL_DB_CONNECTION_STRING
AZURE_STORAGE_CONTAINER_ENVIRONMENT_VARIABLE
SAS_TOKEN_ENVIRONMENT_VARIABLE
SQL_COPT_SS_ACCESS_TOKEN
TOKEN_URL
memory_embedding
results_path
storage_io

AZURE_SQL_DB_CONNECTION_STRING = 'AZURE_SQL_DB_CONNECTION_STRING'

AZURE_STORAGE_CONTAINER_ENVIRONMENT_VARIABLE: str = 'AZURE_STORAGE_ACCOUNT_RESULTS_CONTAINER_URL'

SAS_TOKEN_ENVIRONMENT_VARIABLE: str = 'AZURE_STORAGE_ACCOUNT_RESULTS_SAS_TOKEN'

SQL_COPT_SS_ACCESS_TOKEN = 1256

TOKEN_URL = 'https://database.windows.net/.default'

add_request_pieces_to_memory(*, request_pieces: Sequence[PromptRequestPiece]) → None

Inserts a list of prompt request pieces into the memory storage.

dispose_engine()

Dispose the engine and clean up resources.

get_all_embeddings() → list[EmbeddingDataEntry]

Fetches all entries from the specified table and returns them as model instances.

get_all_prompt_pieces() → list[PromptRequestPiece]

Fetches all entries from the specified table and returns them as model instances.

get_prompt_request_piece_by_memory_labels(*, memory_labels: dict[str, str] = {}) → list[PromptRequestPiece]

Retrieves a list of PromptRequestPiece objects that have the specified memory labels.

Parameters:

• memory_labels (dict[str, str]) – A free-form dictionary for tagging prompts with custom labels.

• operation (These labels can be used to track all prompts sent as part of an)

• on (score prompts based)

• ID (the operation)

• dictionary. (Users can define any key-value pairs according to their needs. Defaults to an empty)

Returns:

A list of PromptRequestPiece with the specified memory labels.

Return type:

list[PromptRequestPiece]

get_prompt_request_pieces_by_id(*, prompt_ids: list[str]) → list[PromptRequestPiece]

Retrieves a list of PromptRequestPiece objects that have the specified prompt ids.

Parameters:

prompt_ids (list[str]) – The prompt IDs to match.

Returns:

A list of PromptRequestPiece with the specified conversation ID.

Return type:

list[PromptRequestPiece]

get_scores_by_prompt_ids(*, prompt_request_response_ids: list[str]) → list[Score]

Gets a list of scores based on prompt_request_response_ids.

get_session() → Session

Provides a session for database operations.

insert_entries(*, entries: list[Base]) → None

Inserts multiple entries into the database.

insert_entry(entry: Base) → None

Inserts an entry into the Table.

Parameters:

entry – An instance of a SQLAlchemy model to be added to the Table.

print_schema()

Prints the schema of all tables in the Azure SQL database.

query_entries(model, *, conditions: Optional = None, distinct: bool = False) → list[Base]

Fetches data from the specified table model with optional conditions.

Parameters:

• model – The SQLAlchemy model class corresponding to the table you want to query.

• conditions – SQLAlchemy filter conditions (Optional).

• distinct – Flag to return distinct rows (defaults to False).

Returns:

List of model instances representing the rows fetched from the table.

reset_database()

Drop and recreate existing tables

update_entries(*, entries: MutableSequence[Base], update_fields: dict) → bool

Updates the given entries with the specified field values.

Parameters:

• entries (list[Base]) – A list of SQLAlchemy model instances to be updated.

• update_fields (dict) – A dictionary of field names and their new values.

Returns:

True if the update was successful, False otherwise.

Return type:

bool