diff --git a/.env_example b/.env_example index 0d70a48cfd..786f0701f2 100644 --- a/.env_example +++ b/.env_example @@ -135,6 +135,10 @@ OPEN_ROUTER_CLAUDE_MODEL="anthropic/claude-3.7-sonnet" OLLAMA_CHAT_ENDPOINT="http://127.0.0.1:11434/v1" OLLAMA_MODEL="llama2" +ABLIT_ENDPOINT="https://api.abliteration.ai/v1" +ABLIT_KEY="ak_xxxxx" +ABLIT_MODEL="abliterated-model" + DEFAULT_OPENAI_FRONTEND_ENDPOINT = ${AZURE_OPENAI_GPT4O_AAD_ENDPOINT} DEFAULT_OPENAI_FRONTEND_KEY = ${AZURE_OPENAI_GPT4O_AAD_KEY} DEFAULT_OPENAI_FRONTEND_MODEL = "gpt-4o" diff --git a/doc/getting_started/populating_secrets.md b/doc/getting_started/populating_secrets.md index 6c784bfef8..110d042291 100644 --- a/doc/getting_started/populating_secrets.md +++ b/doc/getting_started/populating_secrets.md @@ -83,6 +83,16 @@ OPENAI_CHAT_MODEL="anthropic/claude-3.7-sonnet" Get your API key from [openrouter.ai](https://openrouter.ai/). ::: +:::{tab-item} Abliteration.ai +```bash +OPENAI_CHAT_ENDPOINT="https://api.abliteration.ai/v1" +OPENAI_CHAT_KEY="ak_your-key-here" +OPENAI_CHAT_MODEL="abliterated-model" +``` + +Get your API key from [abliteration.ai](https://abliteration.ai/). See [docs.abliteration.ai](https://docs.abliteration.ai) for available models. +::: + :::: ```{note} diff --git a/pyrit/setup/initializers/targets.py b/pyrit/setup/initializers/targets.py index 5b4c700917..aea394d1e4 100644 --- a/pyrit/setup/initializers/targets.py +++ b/pyrit/setup/initializers/targets.py @@ -279,6 +279,13 @@ class TargetConfig: key_var="GOOGLE_GEMINI_API_KEY", model_var="GOOGLE_GEMINI_MODEL", ), + TargetConfig( + registry_name="abliteration", + target_class=OpenAIChatTarget, + endpoint_var="ABLIT_ENDPOINT", + key_var="ABLIT_KEY", + model_var="ABLIT_MODEL", + ), # ============================================ # OpenAI Responses Targets (OpenAIResponseTarget) # ============================================ @@ -499,6 +506,7 @@ class TargetInitializer(PyRITInitializer): - OPEN_ROUTER_* - OpenRouter API - OLLAMA_* - Ollama local - GOOGLE_GEMINI_* - Google Gemini (OpenAI-compatible) + - ABLIT_* - Abliteration.ai (OpenAI-compatible) **OpenAI Responses Targets (OpenAIResponseTarget):** - AZURE_OPENAI_GPT5_RESPONSES_* - Azure OpenAI GPT-5 Responses diff --git a/tests/integration/targets/test_targets_and_secrets.py b/tests/integration/targets/test_targets_and_secrets.py index e2ec9da733..882a09a4db 100644 --- a/tests/integration/targets/test_targets_and_secrets.py +++ b/tests/integration/targets/test_targets_and_secrets.py @@ -1072,6 +1072,7 @@ async def test_openai_chat_target_with_sync_token_provider(sqlite_instance: SQLi ("GROQ_ENDPOINT", "GROQ_KEY", "GROQ_LLAMA_MODEL"), ("OPEN_ROUTER_ENDPOINT", "OPEN_ROUTER_KEY", "OPEN_ROUTER_CLAUDE_MODEL"), ("OLLAMA_CHAT_ENDPOINT", "", "OLLAMA_MODEL"), + ("ABLIT_ENDPOINT", "ABLIT_KEY", "ABLIT_MODEL"), ], ) @pytest.mark.run_only_if_all_tests diff --git a/tests/unit/setup/test_targets_initializer.py b/tests/unit/setup/test_targets_initializer.py index 52141904e1..360bdec747 100644 --- a/tests/unit/setup/test_targets_initializer.py +++ b/tests/unit/setup/test_targets_initializer.py @@ -157,6 +157,22 @@ async def test_registers_ollama_without_api_key(self): assert target is not None assert target._model_name == "llama2" + async def test_registers_abliteration_with_env_vars(self): + """Test that the Abliteration.ai target is registered when its env vars are set.""" + os.environ["ABLIT_ENDPOINT"] = "https://api.abliteration.ai/v1" + os.environ["ABLIT_KEY"] = "test_key" + os.environ["ABLIT_MODEL"] = "test-model" + + init = TargetInitializer() + await init.initialize_async() + + registry = TargetRegistry.get_registry_singleton() + assert "abliteration" in registry.instances + target = registry.instances.get("abliteration") + assert target is not None + assert target._model_name == "test-model" + assert target._endpoint == "https://api.abliteration.ai/v1" + async def test_azure_target_uses_entra_auth(self): """Test that Azure targets use Entra auth (token provider) instead of API key.""" os.environ["AZURE_OPENAI_GPT4O_ENDPOINT"] = "https://my-deployment.openai.azure.com/openai/v1" @@ -205,6 +221,7 @@ def test_expected_targets_in_configs(self): assert "ollama" in registry_names assert "groq" in registry_names assert "google_gemini" in registry_names + assert "abliteration" in registry_names def test_target_configs_have_unique_registry_names(self): """Guard against typos: every ``registry_name`` in ``ENV_TARGET_CONFIGS`` must be unique.