diff --git a/.pyrit_conf_example b/.pyrit_conf_example index 523ccc28e1..1a9fccbf53 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -107,6 +107,12 @@ operation: op_trash_panda # env_akv_ref: # - https://my-vault.vault.azure.net/secrets/my-pyrit-env +# Target API Key Vault +# -------------------- +# Vault used by the backend to store API keys for database-backed targets. +# This is separate from env_akv_ref: it stores one secret per persisted target. +# target_api_key_vault_url: https://my-vault.vault.azure.net + # Max Concurrent Scenario Runs # ---------------------------- # Maximum number of scenario runs that can execute concurrently in the backend. diff --git a/pyrit/auth/key_vault_secret_store.py b/pyrit/auth/key_vault_secret_store.py new file mode 100644 index 0000000000..aa5e840445 --- /dev/null +++ b/pyrit/auth/key_vault_secret_store.py @@ -0,0 +1,80 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Azure Key Vault storage for persisted target API keys.""" + +from urllib.parse import urlparse + + +class KeyVaultSecretStore: + """Store and retrieve target API keys using Azure Key Vault.""" + + def __init__(self, *, vault_url: str) -> None: + """ + Initialize the secret store. + + Args: + vault_url (str): Base URL of the Azure Key Vault. + """ + self._vault_url = vault_url.rstrip("/") + + async def set_secret_async(self, *, name: str, value: str) -> str: + """ + Store a secret and return its versioned Key Vault URI. + + Args: + name (str): Key Vault secret name. + value (str): Secret value to store. + + Returns: + str: The versioned URI returned by Azure Key Vault. + + Raises: + RuntimeError: If Azure Key Vault does not return a secret URI. + """ + from azure.identity.aio import DefaultAzureCredential + from azure.keyvault.secrets.aio import SecretClient + + credential = DefaultAzureCredential() + try: + async with SecretClient(vault_url=self._vault_url, credential=credential) as client: + secret = await client.set_secret(name, value) + finally: + await credential.close() + if not secret.id: + raise RuntimeError(f"Azure Key Vault did not return a URI for secret '{name}'.") + return secret.id + + @staticmethod + async def get_secret_async(*, secret_uri: str) -> str: + """ + Retrieve a secret value from its versioned Key Vault URI. + + Args: + secret_uri (str): A versioned or unversioned Azure Key Vault secret URI. + + Returns: + str: The stored secret value. + + Raises: + ValueError: If the URI is invalid or the secret has no value. + """ + from azure.identity.aio import DefaultAzureCredential + from azure.keyvault.secrets.aio import SecretClient + + parsed = urlparse(secret_uri) + path_parts = [part for part in parsed.path.split("/") if part] + if parsed.scheme != "https" or not parsed.netloc or len(path_parts) not in (2, 3) or path_parts[0] != "secrets": + raise ValueError(f"Invalid Azure Key Vault secret URI: '{secret_uri}'.") + + credential = DefaultAzureCredential() + try: + async with SecretClient( + vault_url=f"{parsed.scheme}://{parsed.netloc}", credential=credential + ) as client: + secret = await client.get_secret(path_parts[1], version=path_parts[2] if len(path_parts) == 3 else None) + finally: + await credential.close() + if secret.value is None: + raise ValueError(f"Azure Key Vault secret '{secret_uri}' has no value.") + return secret.value diff --git a/pyrit/backend/main.py b/pyrit/backend/main.py index 1eb508cf90..41defac9ff 100644 --- a/pyrit/backend/main.py +++ b/pyrit/backend/main.py @@ -34,6 +34,7 @@ targets, version, ) +from pyrit.backend.services.target_service import configure_target_service from pyrit.setup.configuration_loader import ConfigurationLoader # Check for development mode from environment variable @@ -57,6 +58,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: config = ConfigurationLoader.load_with_overrides(config_file=config_file) await config.initialize_pyrit_async() + configure_target_service(api_key_vault_url=config.target_api_key_vault_url) # Expose config values to route handlers via app.state default_labels: dict[str, str] = {} diff --git a/pyrit/backend/services/target_service.py b/pyrit/backend/services/target_service.py index 1743e64145..f28e5bccaa 100644 --- a/pyrit/backend/services/target_service.py +++ b/pyrit/backend/services/target_service.py @@ -12,10 +12,13 @@ - Retrieved from registry (pre-registered at startup or created earlier) """ +import asyncio import logging +import re from functools import lru_cache from typing import Any, Literal, cast +from pyrit.auth.key_vault_secret_store import KeyVaultSecretStore from pyrit.backend.mappers.target_mappers import target_object_to_instance from pyrit.backend.models.common import PaginationInfo from pyrit.backend.models.targets import ( @@ -24,11 +27,15 @@ TargetCatalogResponse, TargetListResponse, ) +from pyrit.memory import CentralMemory +from pyrit.models import OpenAITargetConfig from pyrit.models.catalog.target import TargetInstance from pyrit.registry import TargetRegistry logger = logging.getLogger(__name__) +_target_api_key_vault_url: str | None = None + class TargetService: """ @@ -40,9 +47,10 @@ class TargetService: service only orchestrates the request → registry hand-off. """ - def __init__(self) -> None: + def __init__(self, *, api_key_vault_url: str | None = None) -> None: """Initialize the target service.""" self._registry = TargetRegistry.get_registry_singleton() + self._secret_store = KeyVaultSecretStore(vault_url=api_key_vault_url) if api_key_vault_url else None def _build_instance_from_object(self, *, target_registry_name: str, target_obj: Any) -> TargetInstance: """ @@ -182,7 +190,6 @@ async def create_target_async(self, *, request: CreateTargetRequest) -> TargetIn raise ValueError( f"Target type '{request.type}' not found. Available types: {self._registry.get_class_names()}" ) - target_cls = self._registry.get_class(request.type) params: dict[str, Any] = dict(request.params) @@ -199,11 +206,38 @@ async def create_target_async(self, *, request: CreateTargetRequest) -> TargetIn else: target_obj = target_cls(**params) - self._registry.instances.register(target_obj) - target_registry_name = target_obj.get_identifier().unique_name + if request.type != "OpenAIChatTarget": + self._registry.instances.register(target_obj, name=target_registry_name) + return self._build_instance_from_object(target_registry_name=target_registry_name, target_obj=target_obj) + + secret_uri = await self._persist_api_key_async( + target_registry_name=target_registry_name, + api_key=cast("str | None", params.pop("api_key", None)), + ) + persisted_target = OpenAITargetConfig( + target_registry_name=target_registry_name, + endpoint=cast("str", params["endpoint"]), + model_name=cast("str", params["model_name"]), + auth_mode=request.auth_mode, + api_key_secret_uri=secret_uri, + ) + memory = CentralMemory.get_memory_instance() + await asyncio.to_thread(memory.add_openai_target_config, target=persisted_target) + self._registry.instances.register(target_obj, name=target_registry_name) + return self._build_instance_from_object(target_registry_name=target_registry_name, target_obj=target_obj) + async def _persist_api_key_async(self, *, target_registry_name: str, api_key: str | None) -> str | None: + if not api_key: + return None + if self._secret_store is None: + raise ValueError( + "Target API key persistence requires 'target_api_key_vault_url' in the PyRIT configuration." + ) + secret_name = re.sub(r"[^0-9A-Za-z-]", "-", f"pyrit-target-{target_registry_name}")[:127] + return await self._secret_store.set_secret_async(name=secret_name, value=api_key) + def _has_reference_params(self, *, target_type: str) -> bool: """ Return True if the target type's build contract references other registry @@ -229,4 +263,11 @@ def get_target_service() -> TargetService: Returns: The singleton TargetService instance. """ - return TargetService() + return TargetService(api_key_vault_url=_target_api_key_vault_url) + + +def configure_target_service(*, api_key_vault_url: str | None) -> None: + """Configure the cached backend target service.""" + global _target_api_key_vault_url + _target_api_key_vault_url = api_key_vault_url + get_target_service.cache_clear() diff --git a/pyrit/memory/alembic/versions/4a7c9e1b3d5f_add_targets_table.py b/pyrit/memory/alembic/versions/4a7c9e1b3d5f_add_targets_table.py new file mode 100644 index 0000000000..19254ff505 --- /dev/null +++ b/pyrit/memory/alembic/versions/4a7c9e1b3d5f_add_targets_table.py @@ -0,0 +1,39 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Add persisted target configurations. + +Revision ID: 4a7c9e1b3d5f +Revises: 3f6e8a0c2d4b +Create Date: 2026-07-17 00:00:00.000000 +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "4a7c9e1b3d5f" +down_revision: str | None = "3f6e8a0c2d4b" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + """Create the Targets table.""" + op.create_table( + "Targets", + sa.Column("target_registry_name", sa.String(), nullable=False), + sa.Column("endpoint", sa.String(), nullable=False), + sa.Column("model_name", sa.String(), nullable=False), + sa.Column("auth_mode", sa.String(), nullable=False), + sa.Column("api_key_secret_uri", sa.String(), nullable=True), + sa.PrimaryKeyConstraint("target_registry_name"), + ) + + +def downgrade() -> None: + """Drop the Targets table.""" + op.drop_table("Targets") diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 602f13b6bf..40df671465 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -32,6 +32,7 @@ ConversationEntry, ConverterIdentifierEntry, EmbeddingDataEntry, + OpenAITargetConfigEntry, PromptConverterIdentifierEntry, PromptMemoryEntry, ScenarioIdentifierEntry, @@ -63,6 +64,7 @@ IdentifierType, Message, MessagePiece, + OpenAITargetConfig, ScenarioIdentifier, ScenarioResult, Score, @@ -149,6 +151,15 @@ def enable_embedding(self, embedding_model: Any | None = None) -> None: self.memory_embedding = default_memory_embedding_factory(embedding_model=embedding_model) + def add_openai_target_config(self, *, target: OpenAITargetConfig) -> None: + """Persist a sanitized, reconstructable target configuration.""" + self._insert_entry(OpenAITargetConfigEntry.from_domain_model(target)) + + def get_openai_target_configs(self) -> Sequence[OpenAITargetConfig]: + """Return all persisted target configurations.""" + entries = self._query_entries(OpenAITargetConfigEntry, order_by=OpenAITargetConfigEntry.target_registry_name) + return [entry.to_domain_model() for entry in entries] + def disable_embedding(self) -> None: """ Disable embedding functionality for the memory interface. diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 24e8fb875c..3e938bfca5 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -53,6 +53,7 @@ ConverterIdentifier, EvaluationIdentifier, MessagePiece, + OpenAITargetConfig, PromptDataType, ScenarioEvaluationIdentifier, ScenarioIdentifier, @@ -419,6 +420,47 @@ def __init_subclass__(cls, **kwargs: Any) -> None: ) +class OpenAITargetConfigEntry(DomainBackedEntry[OpenAITargetConfig]): + """Persistence projection for a reconstructable target configuration.""" + + __tablename__ = "Targets" + __table_args__ = {"extend_existing": True} + + target_registry_name: Mapped[str] = mapped_column(String, primary_key=True) + endpoint: Mapped[str] = mapped_column(String, nullable=False) + model_name: Mapped[str] = mapped_column(String, nullable=False) + auth_mode: Mapped[Literal["api_key", "identity"]] = mapped_column(String, nullable=False) + api_key_secret_uri: Mapped[str | None] = mapped_column(String, nullable=True) + + @classmethod + def from_domain_model(cls, domain_model: OpenAITargetConfig) -> Self: + """ + Build an unsaved target row from its domain model. + + Args: + domain_model (OpenAITargetConfig): The target configuration to persist. + + Returns: + Self: An unsaved target configuration row. + """ + return cls(**domain_model.model_dump()) + + def to_domain_model(self) -> OpenAITargetConfig: + """ + Build the canonical domain model represented by this row. + + Returns: + OpenAITargetConfig: The reconstructed target configuration. + """ + return OpenAITargetConfig( + target_registry_name=self.target_registry_name, + endpoint=self.endpoint, + model_name=self.model_name, + auth_mode=self.auth_mode, + api_key_secret_uri=self.api_key_secret_uri, + ) + + T = TypeVar("T", bound=ComponentIdentifier) diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 4c9075d58a..664699480c 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -83,6 +83,7 @@ RegistryReference, display_choices, ) +from pyrit.models.openai_target_config import OpenAITargetConfig from pyrit.models.question_answering import QuestionAnsweringDataset, QuestionAnsweringEntry, QuestionChoice from pyrit.models.results.attack_result import AttackOutcome, AttackResult, AttackResultT from pyrit.models.results.scenario_result import ScenarioResult, ScenarioRunState @@ -181,6 +182,7 @@ "ObjectiveTargetEvaluationIdentifier", "Parameter", "ParameterDestination", + "OpenAITargetConfig", "PromptDataType", "PromptResponseError", "QuestionAnsweringDataset", diff --git a/pyrit/models/openai_target_config.py b/pyrit/models/openai_target_config.py new file mode 100644 index 0000000000..5af7240fa2 --- /dev/null +++ b/pyrit/models/openai_target_config.py @@ -0,0 +1,26 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Domain model for OpenAI target configurations persisted in memory.""" + +from typing import Literal + +from pydantic import BaseModel, model_validator + + +class OpenAITargetConfig(BaseModel): + """A reconstructable OpenAI target configuration with a secret reference.""" + + target_registry_name: str + endpoint: str + model_name: str + auth_mode: Literal["api_key", "identity"] = "api_key" + api_key_secret_uri: str | None = None + + @model_validator(mode="after") + def _validate_secret_storage(self) -> "OpenAITargetConfig": + if self.auth_mode == "api_key" and not self.api_key_secret_uri: + raise ValueError("API key authentication requires an api_key_secret_uri.") + if self.auth_mode == "identity" and self.api_key_secret_uri: + raise ValueError("Identity authentication must not reference an API key secret.") + return self diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 4ed9355deb..5ea2bf6293 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -122,6 +122,7 @@ class ConfigurationLoader(YamlLoadable): operation: str | None = None max_concurrent_scenario_runs: int = 3 allow_custom_initializers: bool = False + target_api_key_vault_url: str | None = None server: dict[str, Any] | None = None extensions: dict[str, Any] = field(default_factory=dict) @@ -316,6 +317,7 @@ def load_with_overrides( config_data["initialization_scripts"] = default_config.initialization_scripts config_data["env_files"] = default_config.env_files config_data["env_akv_ref"] = default_config.env_akv_ref + config_data["target_api_key_vault_url"] = default_config.target_api_key_vault_url config_data["silent"] = default_config.silent if default_config.operator: config_data["operator"] = default_config.operator @@ -342,6 +344,7 @@ def load_with_overrides( config_data["initialization_scripts"] = explicit_config.initialization_scripts config_data["env_files"] = explicit_config.env_files config_data["env_akv_ref"] = explicit_config.env_akv_ref + config_data["target_api_key_vault_url"] = explicit_config.target_api_key_vault_url config_data["silent"] = explicit_config.silent if explicit_config.operator: config_data["operator"] = explicit_config.operator diff --git a/pyrit/setup/initializers/dbtargets.py b/pyrit/setup/initializers/dbtargets.py new file mode 100644 index 0000000000..b49fba6dab --- /dev/null +++ b/pyrit/setup/initializers/dbtargets.py @@ -0,0 +1,34 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Initializer for targets persisted in the PyRIT database.""" + +import asyncio + +from pyrit.auth.key_vault_secret_store import KeyVaultSecretStore +from pyrit.memory import CentralMemory +from pyrit.prompt_target import OpenAIChatTarget +from pyrit.registry import TargetRegistry +from pyrit.setup.pyrit_initializer import PyRITInitializer + + +class DbtargetsInitializer(PyRITInitializer): + """Load database-backed OpenAI chat targets into the target registry.""" + + async def initialize_async(self) -> None: + """Load persisted target configurations and register their target instances.""" + memory = CentralMemory.get_memory_instance() + targets = await asyncio.to_thread(memory.get_openai_target_configs) + registry = TargetRegistry.get_registry_singleton() + + for target in targets: + api_key = None + if target.api_key_secret_uri: + api_key = await KeyVaultSecretStore.get_secret_async(secret_uri=target.api_key_secret_uri) + + target_instance = OpenAIChatTarget( + endpoint=target.endpoint, + model_name=target.model_name, + api_key=api_key, + ) + registry.instances.register(target_instance, name=target.target_registry_name) diff --git a/tests/unit/backend/test_target_service.py b/tests/unit/backend/test_target_service.py index 8a70b7c9bc..d8408e2ce9 100644 --- a/tests/unit/backend/test_target_service.py +++ b/tests/unit/backend/test_target_service.py @@ -6,7 +6,7 @@ """ import os -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -301,7 +301,7 @@ async def test_create_target_registers_in_registry(self, sqlite_instance) -> Non async def test_create_target_model_name_not_overridden_by_env_var(self, sqlite_instance) -> None: """Test that explicit model_name is not overridden by underlying_model env var.""" with patch.dict(os.environ, {"OPENAI_CHAT_UNDERLYING_MODEL": "gpt-4o"}): - service = TargetService() + service = TargetService(api_key_vault_url="https://vault.vault.azure.net") request = CreateTargetRequest( type="OpenAIChatTarget", @@ -312,7 +312,13 @@ async def test_create_target_model_name_not_overridden_by_env_var(self, sqlite_i }, ) - result = await service.create_target_async(request=request) + with patch.object( + service._secret_store, + "set_secret_async", + new_callable=AsyncMock, + return_value="https://vault.vault.azure.net/secrets/key/version", + ): + result = await service.create_target_async(request=request) assert result.identifier.model_name == "claude-sonnet-4-6" # underlying_model_name is empty since no underlying_model was passed @@ -320,7 +326,7 @@ async def test_create_target_model_name_not_overridden_by_env_var(self, sqlite_i async def test_create_target_with_different_underlying_model(self, sqlite_instance) -> None: """Test that explicit underlying_model is used when it differs from model_name.""" - service = TargetService() + service = TargetService(api_key_vault_url="https://vault.vault.azure.net") request = CreateTargetRequest( type="OpenAIChatTarget", @@ -332,11 +338,85 @@ async def test_create_target_with_different_underlying_model(self, sqlite_instan }, ) - result = await service.create_target_async(request=request) + with patch.object( + service._secret_store, + "set_secret_async", + new_callable=AsyncMock, + return_value="https://vault.vault.azure.net/secrets/key/version", + ): + result = await service.create_target_async(request=request) assert result.identifier.model_name == "my-gpt4o-deployment" assert result.identifier.underlying_model_name == "gpt-4o" + async def test_create_api_key_target_requires_configured_vault(self, sqlite_instance) -> None: + service = TargetService() + request = CreateTargetRequest( + type="OpenAIChatTarget", + params={ + "model_name": "model", + "endpoint": "https://example.test", + "api_key": "secret", + }, + ) + + with pytest.raises(ValueError, match="target_api_key_vault_url"): + await service.create_target_async(request=request) + + assert sqlite_instance.get_openai_target_configs() == [] + assert len(service._registry.instances) == 0 + + async def test_create_api_key_target_persists_secret_reference(self, sqlite_instance) -> None: + service = TargetService(api_key_vault_url="https://vault.vault.azure.net") + request = CreateTargetRequest( + type="OpenAIChatTarget", + params={ + "model_name": "model", + "endpoint": "https://example.test", + "api_key": "secret", + }, + ) + secret_uri = "https://vault.vault.azure.net/secrets/key/version" + + with patch.object( + service._secret_store, + "set_secret_async", + new_callable=AsyncMock, + return_value=secret_uri, + ): + result = await service.create_target_async(request=request) + + configs = sqlite_instance.get_openai_target_configs() + assert len(configs) == 1 + assert configs[0].target_registry_name == result.target_registry_name + assert configs[0].api_key_secret_uri == secret_uri + assert "secret" not in configs[0].model_dump().values() + + async def test_create_api_key_target_does_not_persist_when_vault_write_fails(self, sqlite_instance) -> None: + service = TargetService(api_key_vault_url="https://vault.vault.azure.net") + request = CreateTargetRequest( + type="OpenAIChatTarget", + params={ + "model_name": "model", + "endpoint": "https://example.test", + "api_key": "secret", + }, + ) + + with ( + patch.object( + service._secret_store, + "set_secret_async", + new_callable=AsyncMock, + side_effect=RuntimeError("vault unavailable"), + ), + pytest.raises(RuntimeError, match="vault unavailable"), + ): + await service.create_target_async(request=request) + + assert sqlite_instance.get_openai_target_configs() == [] + assert len(service._registry.instances) == 0 + class TestCreateTargetEntraAuth: """Entra auth at the service boundary: the service only omits the api_key and diff --git a/tests/unit/memory/test_openai_target_config.py b/tests/unit/memory/test_openai_target_config.py new file mode 100644 index 0000000000..07d24f0d61 --- /dev/null +++ b/tests/unit/memory/test_openai_target_config.py @@ -0,0 +1,17 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from pyrit.models import OpenAITargetConfig + + +def test_openai_target_config_round_trip(sqlite_instance) -> None: + config = OpenAITargetConfig( + target_registry_name="saved-target", + endpoint="https://example.test", + model_name="model", + auth_mode="identity", + ) + + sqlite_instance.add_openai_target_config(target=config) + + assert sqlite_instance.get_openai_target_configs() == [config] diff --git a/tests/unit/models/test_openai_target_config.py b/tests/unit/models/test_openai_target_config.py new file mode 100644 index 0000000000..632b98c1c7 --- /dev/null +++ b/tests/unit/models/test_openai_target_config.py @@ -0,0 +1,28 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import pytest +from pydantic import ValidationError + +from pyrit.models import OpenAITargetConfig + + +def test_api_key_auth_requires_secret_uri() -> None: + with pytest.raises(ValidationError, match="api_key_secret_uri"): + OpenAITargetConfig( + target_registry_name="target", + endpoint="https://example.test", + model_name="model", + auth_mode="api_key", + ) + + +def test_identity_auth_rejects_secret_uri() -> None: + with pytest.raises(ValidationError, match="must not reference"): + OpenAITargetConfig( + target_registry_name="target", + endpoint="https://example.test", + model_name="model", + auth_mode="identity", + api_key_secret_uri="https://vault.vault.azure.net/secrets/key/version", + ) diff --git a/tests/unit/setup/test_configuration_loader.py b/tests/unit/setup/test_configuration_loader.py index a42de74afe..61fb2ad6c1 100644 --- a/tests/unit/setup/test_configuration_loader.py +++ b/tests/unit/setup/test_configuration_loader.py @@ -42,8 +42,17 @@ def test_default_values(self): assert config.initialization_scripts is None # None means "use defaults" assert config.env_files is None # None means "use defaults" assert config.env_akv_ref is None + assert config.target_api_key_vault_url is None assert config.silent is False + def test_target_api_key_vault_url_from_dict(self): + """Test loading the dedicated target API key vault URL.""" + config = ConfigurationLoader.from_dict( + {"target_api_key_vault_url": "https://target-secrets.vault.azure.net"} + ) + + assert config.target_api_key_vault_url == "https://target-secrets.vault.azure.net" + def test_valid_memory_db_types_snake_case(self): """Test all valid memory database types in snake_case.""" for db_type in ["in_memory", "sqlite", "azure_sql"]: diff --git a/tests/unit/setup/test_dbtargets_initializer.py b/tests/unit/setup/test_dbtargets_initializer.py new file mode 100644 index 0000000000..cbe0db736b --- /dev/null +++ b/tests/unit/setup/test_dbtargets_initializer.py @@ -0,0 +1,51 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from unittest.mock import AsyncMock, patch + +from pyrit.models import OpenAITargetConfig +from pyrit.registry import TargetRegistry +from pyrit.setup.initializers.dbtargets import DbtargetsInitializer + + +async def test_initialize_registers_api_key_target(sqlite_instance) -> None: + sqlite_instance.add_openai_target_config( + target=OpenAITargetConfig( + target_registry_name="saved-target", + endpoint="https://example.test", + model_name="model", + api_key_secret_uri="https://vault.vault.azure.net/secrets/key/version", + ) + ) + + with patch( + "pyrit.setup.initializers.dbtargets.KeyVaultSecretStore.get_secret_async", + new_callable=AsyncMock, + return_value="secret-value", + ) as get_secret: + await DbtargetsInitializer().initialize_async() + + target = TargetRegistry.get_registry_singleton().instances.get("saved-target") + assert target is not None + assert target._api_key == "secret-value" + get_secret.assert_awaited_once() + + +async def test_initialize_registers_identity_target(sqlite_instance) -> None: + sqlite_instance.add_openai_target_config( + target=OpenAITargetConfig( + target_registry_name="identity-target", + endpoint="https://example.openai.azure.com", + model_name="model", + auth_mode="identity", + ) + ) + + token_provider = AsyncMock(return_value="token") + with patch( + "pyrit.prompt_target.openai.openai_target.get_azure_openai_auth", + return_value=token_provider, + ): + await DbtargetsInitializer().initialize_async() + + assert TargetRegistry.get_registry_singleton().instances.get("identity-target") is not None