From 0252be8640542a716ea680dad2f187b941b8919c Mon Sep 17 00:00:00 2001 From: Zippo00 Date: Sat, 15 Aug 2026 13:21:58 +0300 Subject: [PATCH 1/3] added --device flag for cli to set device used by elm and alm; fixed triton-windows error on Windows --- avise/cli.py | 9 ++++++- avise/engine.py | 24 +++++++++++++------ avise/pipelines/languagemodel/pipeline.py | 5 ++++ .../multi_turn/deceptive_delight.py | 4 ++++ .../languagemodel/multi_turn/red_queen.py | 3 +++ .../single_turn/prompt_injection.py | 2 ++ 6 files changed, 39 insertions(+), 8 deletions(-) diff --git a/avise/cli.py b/avise/cli.py index 4688c8b..792a6f1 100644 --- a/avise/cli.py +++ b/avise/cli.py @@ -145,6 +145,12 @@ def main(arguments=None) -> None: "-a", help="API Key to use with requests sent to target API (overrides api_key from Connector configuration file).", ) + parser.add_argument( + "--device", + default=None, + choices=[None, "auto", "cpu", "gpu"], + help='Which device to load local models on ("auto", "cpu", or "gpu"). If given, overrides the device setting from SET configuration file.', + ) parser.add_argument("--version", "-V", action="version", version=__version__) args = parser.parse_args(arguments) @@ -231,6 +237,7 @@ def main(arguments=None) -> None: output_path=args.output, target=args.target, api_key=args.api_key, + device=args.device ) # Print a small summary to the console @@ -243,7 +250,7 @@ def main(arguments=None) -> None: print( f" Failed: {report.summary['failed']} ({report.summary['fail_rate']}%)" ) - print(f" Errors: {report.summary['error']}") + print(f" Inconclusive: {report.summary['error']}") except Exception as e: logger.error( diff --git a/avise/engine.py b/avise/engine.py index 2ec4111..55d9c17 100644 --- a/avise/engine.py +++ b/avise/engine.py @@ -33,20 +33,28 @@ # On Windows, ensure triton-windows package is installed if os.name == "nt": - if importlib.util.find_spec("triton-windows") is None: - logger.info( - "The current Operating System seems to be Windows. We need to install triton-windows Python package to the current environment in order to run required language models." + if importlib.util.find_spec("triton") is None: + logger.warning( + "triton-windows does not appear to be installed. This should have " + "been installed automatically as a dependency on Windows." ) + try: + import pip # noqa: F401 + except ImportError: + raise RuntimeError( + "triton-windows is missing and pip is not available in this " + "environment to install it automatically (this is common with " + "'uv tool' or 'pipx' installs). Please reinstall the package, " + "or run: uv tool install --with triton-windows" + ) try: subprocess.check_call( [sys.executable, "-m", "pip", "install", "triton-windows"] ) - logger.info( - "Successfully installed triton-windows package to the current environment." - ) except Exception as e: raise RuntimeError( - "Unable to install triton-windows Python package. Cannot run required language models on Windows without it. Try pip install triton-windows" + "Unable to install triton-windows Python package. " + "Try: pip install triton-windows" ) from e @@ -107,6 +115,7 @@ def run_test( output_path: Optional[str] = None, target: Optional[str] = None, api_key: Optional[str] = None, + device: Optional[str] = None ) -> dict: """Run the 4-phase pipeline @@ -183,6 +192,7 @@ def run_test( connector_config_path=connector_config_path, generate_ai_summary=generate_ai_summary, runs=runs, + device=device ) def _build_connector(self, connector_config: dict, evaluation: bool = False) -> Any: diff --git a/avise/pipelines/languagemodel/pipeline.py b/avise/pipelines/languagemodel/pipeline.py index 38e9896..9c9c58a 100644 --- a/avise/pipelines/languagemodel/pipeline.py +++ b/avise/pipelines/languagemodel/pipeline.py @@ -63,6 +63,7 @@ def __init__(self): self.evaluation_model_name: Optional[str] = None self.evaluation_model_max_tokens: Optional[int] = None self.evaluation_model: Optional[EvaluationLanguageModel] = None + self.device: Optional[str] = None @abstractmethod def initialize(self, set_config_path: str) -> List[LanguageModelSETCase]: @@ -150,6 +151,7 @@ def run( connector_config_path: Optional[str] = None, generate_ai_summary: bool = True, runs: int = 1, + device: Optional[str] = None ) -> ReportData: """Orchestration method that executes the 4-phase pipeline. This method gets called by the execution engine. @@ -175,6 +177,9 @@ def run( self.set_config_path = set_config_path self.target_model_name = connector.model + # Store device + self.device = device + try: # Initialize sets = self.initialize(set_config_path) diff --git a/avise/sets/languagemodel/multi_turn/deceptive_delight.py b/avise/sets/languagemodel/multi_turn/deceptive_delight.py index e45930d..3b30095 100644 --- a/avise/sets/languagemodel/multi_turn/deceptive_delight.py +++ b/avise/sets/languagemodel/multi_turn/deceptive_delight.py @@ -46,6 +46,10 @@ def initialize(self, set_config_path: str) -> List[LanguageModelSETCase]: set_config = ConfigLoader().load(set_config_path) + if self.device is not None: + set_config["evaluation_model_device"] = str(self.device) + set_config["adversarial_model_device"] = str(self.device) + sets = set_config.get("sets", []) if not sets: raise ValueError( diff --git a/avise/sets/languagemodel/multi_turn/red_queen.py b/avise/sets/languagemodel/multi_turn/red_queen.py index 04757f6..810d6a3 100644 --- a/avise/sets/languagemodel/multi_turn/red_queen.py +++ b/avise/sets/languagemodel/multi_turn/red_queen.py @@ -49,6 +49,9 @@ def initialize(self, set_config_path: str) -> List[LanguageModelSETCase]: logger.info(f"Initializing Security Evaluation Test: {self.name}") set_config = ConfigLoader().load(set_config_path) + if self.device is not None: + set_config["evaluation_model_device"] = str(self.device) + set_config["adversarial_model_device"] = str(self.device) sets = set_config.get("sets", []) if not sets: diff --git a/avise/sets/languagemodel/single_turn/prompt_injection.py b/avise/sets/languagemodel/single_turn/prompt_injection.py index aed00cf..7422752 100644 --- a/avise/sets/languagemodel/single_turn/prompt_injection.py +++ b/avise/sets/languagemodel/single_turn/prompt_injection.py @@ -74,6 +74,8 @@ def initialize(self, set_config_path: str) -> List[LanguageModelSETCase]: logger.info(f"Initializing SET: {self.name}") config = ConfigLoader().load(set_config_path) + if self.device is not None: + config["evaluation_model_device"] = str(self.device) self.evaluation_system_prompt = config.get("evaluation_system_prompt") if self.evaluation_system_prompt and self.evaluation_model_name: From 39e99c196d28961a168ae87ce97b7a2b58f089ab Mon Sep 17 00:00:00 2001 From: Zippo00 Date: Sat, 15 Aug 2026 13:23:18 +0300 Subject: [PATCH 2/3] version bump --- avise/__init__.py | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/avise/__init__.py b/avise/__init__.py index 090b930..b9421a8 100644 --- a/avise/__init__.py +++ b/avise/__init__.py @@ -1,3 +1,3 @@ -__version__ = "0.2.4" +__version__ = "0.2.5" __app__ = "AVISE" __description__ = "AI Vulnerability Identification & Security Evaluation framework" \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 0003d4c..f82fc04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "hatchling.build" [project] name = "avise" -version = "0.2.4" +version = "0.2.5" authors = [ { name = "Mikko Lempinen", email="mikko.lempinen@oulu.fi" }, { name = "Joni Kemppainen" }, From b8048f8cc600b460b1b17160cbbb2ada54ce8c8d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci-lite[bot]" <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com> Date: Sat, 15 Aug 2026 10:25:23 +0000 Subject: [PATCH 3/3] [pre-commit.ci lite] apply automatic fixes --- avise/cli.py | 2 +- avise/engine.py | 4 ++-- avise/pipelines/languagemodel/pipeline.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/avise/cli.py b/avise/cli.py index 792a6f1..5dddc1c 100644 --- a/avise/cli.py +++ b/avise/cli.py @@ -237,7 +237,7 @@ def main(arguments=None) -> None: output_path=args.output, target=args.target, api_key=args.api_key, - device=args.device + device=args.device, ) # Print a small summary to the console diff --git a/avise/engine.py b/avise/engine.py index 55d9c17..c75451c 100644 --- a/avise/engine.py +++ b/avise/engine.py @@ -115,7 +115,7 @@ def run_test( output_path: Optional[str] = None, target: Optional[str] = None, api_key: Optional[str] = None, - device: Optional[str] = None + device: Optional[str] = None, ) -> dict: """Run the 4-phase pipeline @@ -192,7 +192,7 @@ def run_test( connector_config_path=connector_config_path, generate_ai_summary=generate_ai_summary, runs=runs, - device=device + device=device, ) def _build_connector(self, connector_config: dict, evaluation: bool = False) -> Any: diff --git a/avise/pipelines/languagemodel/pipeline.py b/avise/pipelines/languagemodel/pipeline.py index 9c9c58a..84b5115 100644 --- a/avise/pipelines/languagemodel/pipeline.py +++ b/avise/pipelines/languagemodel/pipeline.py @@ -151,7 +151,7 @@ def run( connector_config_path: Optional[str] = None, generate_ai_summary: bool = True, runs: int = 1, - device: Optional[str] = None + device: Optional[str] = None, ) -> ReportData: """Orchestration method that executes the 4-phase pipeline. This method gets called by the execution engine.