diff --git a/agentic_security/probe_actor/fuzzer.py b/agentic_security/probe_actor/fuzzer.py index 070eb12..802eab4 100644 --- a/agentic_security/probe_actor/fuzzer.py +++ b/agentic_security/probe_actor/fuzzer.py @@ -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, diff --git a/agentic_security/probe_data/modules/garak_tool.py b/agentic_security/probe_data/modules/garak_tool.py index 9feb430..f1d2967 100644 --- a/agentic_security/probe_data/modules/garak_tool.py +++ b/agentic_security/probe_data/modules/garak_tool.py @@ -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", + } + } +} diff --git a/agentic_security/routes/proxy.py b/agentic_security/routes/proxy.py index ac75ac5..45b15fa 100644 --- a/agentic_security/routes/proxy.py +++ b/agentic_security/routes/proxy.py @@ -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) diff --git a/agentic_security/test_lib.py b/agentic_security/test_lib.py index 9e4d95b..881d9ba 100644 --- a/agentic_security/test_lib.py +++ b/agentic_security/test_lib.py @@ -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] diff --git a/poetry.lock b/poetry.lock index a78c5f2..abec27d 100644 --- a/poetry.lock +++ b/poetry.lock @@ -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"},