fix(garak runner):

This commit is contained in:
Alexander Myasoedov
2024-12-23 16:45:44 +02:00
parent f844ed11df
commit 79cbdf6c4d
5 changed files with 71 additions and 11 deletions
+2
View File
@@ -13,6 +13,8 @@ from agentic_security.probe_actor.refusal import refusal_heuristic
from agentic_security.probe_data import msj_data
from agentic_security.probe_data.data import prepare_prompts
# TODO: full log file
async def generate_prompts(
prompts: list[str] | AsyncGenerator,
@@ -1,5 +1,6 @@
import asyncio
import importlib.util
import json
import os
import subprocess
@@ -7,6 +8,16 @@ from loguru import logger
# TODO: add probes modules
GARAK_CONFIG = "garak_rest.json"
def write_garak_config_json():
with open(GARAK_CONFIG, "w") as f:
f.write(json.dumps(SPEC))
# TODO: add config params to data registry
class Module:
def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue):
@@ -22,20 +33,20 @@ class Module:
async def apply(self) -> []:
env = os.environ.copy()
env["OPENAI_API_BASE"] = "http://0.0.0.0:8718/proxy"
# Command to be executed
command = [
"python",
"-m",
"garak",
"--model_type",
"openai",
"--model_name",
"gpt-3.5-turbo",
"rest",
"-G",
GARAK_CONFIG,
"--probes",
"encoding",
]
logger.info("Starting Garak tool. Writing config file.")
write_garak_config_json()
logger.info(f"Executing command: {command}")
# Execute the command with the specific environment
process = subprocess.Popen(
@@ -57,3 +68,26 @@ class Module:
logger.error(f"Error executing command: {command}")
logger.error(f"err: {err}")
return
SPEC = {
"rest": {
"RestGenerator": {
"name": "OpenAI GPT-4 Service",
"uri": "http://0.0.0.0:8718/proxy/chat/completions",
"method": "POST",
"headers": {
"Authorization": "Bearer $OPENAI_API_KEY",
"Content-Type": "application/json",
},
"req_template_json_object": {
"model": "gpt-4",
"messages": [{"role": "user", "content": "$INPUT"}],
"max_tokens": 1050,
"temperature": 0.7,
},
"response_json": True,
"response_json_field": "$.choices[0].message.content",
}
}
}
+1
View File
@@ -17,6 +17,7 @@ async def proxy_completions(request: CompletionRequest):
prompt_content = " ".join(
[msg.content for msg in request.messages if msg.role == "user"]
)
# Todo: get current llm spec for proper proxing
message = prompt_content + " " + message
ready = Event()
ref = dict(message=message, reply="", ready=ready)
+28
View File
@@ -1,3 +1,4 @@
import importlib
import signal
import subprocess
import time
@@ -8,6 +9,11 @@ import agentic_security.test_spec_assets as test_spec_assets
from agentic_security.lib import AgenticSecurity
def has_module(module_name):
module_obj = importlib.util.find_spec(module_name)
return module_obj is not None
@pytest.fixture(scope="module")
def test_server(request):
# Start server process
@@ -66,3 +72,25 @@ class TestAS:
assert isinstance(result, dict)
print(result)
assert len(result) in [0, 1]
@pytest.mark.skipif(not has_module("garak"), reason="Garak module not installed")
def test_garak(self, test_server):
llmSpec = test_spec_assets.SAMPLE_SPEC
maxBudget = 1000000
max_th = 0.3
datasets = [
{
"dataset_name": "Garak",
"num_prompts": 10,
"tokens": 0,
"approx_cost": 0.0,
"source": "Github: https://github.com/leondz/garak#v0.9.0.1",
"selected": True,
"url": "https://github.com/leondz/garak2",
"dynamic": True,
},
]
result = AgenticSecurity.scan(llmSpec, maxBudget, datasets, max_th)
assert isinstance(result, dict)
print(result)
assert len(result) in [0, 1]
Generated
+1 -6
View File
@@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 1.8.5 and should not be changed by hand.
# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand.
[[package]]
name = "aiohappyeyeballs"
@@ -2520,11 +2520,6 @@ files = [
{file = "scikit_learn-1.5.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f60021ec1574e56632be2a36b946f8143bf4e5e6af4a06d85281adc22938e0dd"},
{file = "scikit_learn-1.5.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:394397841449853c2290a32050382edaec3da89e35b3e03d6cc966aebc6a8ae6"},
{file = "scikit_learn-1.5.2-cp312-cp312-win_amd64.whl", hash = "sha256:57cc1786cfd6bd118220a92ede80270132aa353647684efa385a74244a41e3b1"},
{file = "scikit_learn-1.5.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e9a702e2de732bbb20d3bad29ebd77fc05a6b427dc49964300340e4c9328b3f5"},
{file = "scikit_learn-1.5.2-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:b0768ad641981f5d3a198430a1d31c3e044ed2e8a6f22166b4d546a5116d7908"},
{file = "scikit_learn-1.5.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:178ddd0a5cb0044464fc1bfc4cca5b1833bfc7bb022d70b05db8530da4bb3dd3"},
{file = "scikit_learn-1.5.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f7284ade780084d94505632241bf78c44ab3b6f1e8ccab3d2af58e0e950f9c12"},
{file = "scikit_learn-1.5.2-cp313-cp313-win_amd64.whl", hash = "sha256:b7b0f9a0b1040830d38c39b91b3a44e1b643f4b36e36567b80b7c6bd2202a27f"},
{file = "scikit_learn-1.5.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:757c7d514ddb00ae249832fe87100d9c73c6ea91423802872d9e74970a0e40b9"},
{file = "scikit_learn-1.5.2-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:52788f48b5d8bca5c0736c175fa6bdaab2ef00a8f536cda698db61bd89c551c1"},
{file = "scikit_learn-1.5.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:643964678f4b5fbdc95cbf8aec638acc7aa70f5f79ee2cdad1eec3df4ba6ead8"},