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, results_container_url: str | None = None, results_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_conversations(*[, orchestrator_id, ...])	Exports conversation data with the given inputs to a specified file.
get_all_embeddings()	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_request_pieces(*[, ...])	Retrieves a list of PromptRequestPiece objects based on the specified filters.
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.
populate_prompt_piece_scores(...)	Adds scores in the database to prompt request piece objects
print_schema()	Prints the schema of all tables in the Azure SQL database.
reset_database()	Drop and recreate existing tables
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_ACCOUNT_RESULTS_CONTAINER_URL
AZURE_STORAGE_ACCOUNT_RESULTS_SAS_TOKEN
SQL_COPT_SS_ACCESS_TOKEN
TOKEN_URL
memory_embedding
results_path
results_storage_io

AZURE_SQL_DB_CONNECTION_STRING = 'AZURE_SQL_DB_CONNECTION_STRING'

AZURE_STORAGE_ACCOUNT_RESULTS_CONTAINER_URL: str = 'AZURE_STORAGE_ACCOUNT_RESULTS_CONTAINER_URL'

AZURE_STORAGE_ACCOUNT_RESULTS_SAS_TOKEN: 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_session() → Session

Provides a session for database operations.

print_schema()

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

reset_database()

Drop and recreate existing tables