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/avise/cli.py b/avise/cli.py index 4688c8b..5dddc1c 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..c75451c 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..84b5115 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: 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" },