diff --git a/agentic_security/probe_data/modules/inspect_ai_tool.py b/agentic_security/probe_data/modules/inspect_ai_tool.py index 78afea8..1938f7c 100644 --- a/agentic_security/probe_data/modules/inspect_ai_tool.py +++ b/agentic_security/probe_data/modules/inspect_ai_tool.py @@ -1,83 +1,83 @@ -import asyncio -import importlib.util -import os - -from agentic_security.logutils import logger - -inspect_ai_task = ( - __file__.replace("inspect_ai_tool.py", "inspect_ai_task.py") - .replace(os.getcwd(), "") - .strip("/") -) - - -class Module: - """:class:`Module` that runs Inspect AI evaluations against an LLM proxy. - - Launches the `Inspect AI `_ - evaluation framework via ``inspect eval``, targeting a local LLM proxy. - The module requires ``inspect_ai`` to be installed - (``pip install inspect_ai``). - - Configuration is passed via ``opts``: - - * ``port`` (int): LLM proxy base URL port. Defaults to ``8718``. - - Attributes: - tools_inbox: Async queue that receives evaluation results. - opts: Configuration dictionary. - """ - - name = "Inspect AI" - - def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue, opts: dict = {}): - self.tools_inbox = tools_inbox - if not self.is_tool_installed(): - logger.error( - "inspect_ai module is not installed. Please install it using '"'"'pip install inspect_ai'"'"'" - ) - self.opts = opts - - def is_tool_installed(self) -> bool: - """Return ``True`` if the ``inspect_ai`` package is importable.""" - inspect_ai = importlib.util.find_spec("inspect_ai") - return inspect_ai is not None - - async def _proc(self, command): - env = os.environ.copy() - process = await asyncio.create_subprocess_shell( - command, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - env=env, - shell=True, - ) - - logger.info(f"Started {command}") - - async for line in process.stdout: - logger.info(line.decode().strip()) - - err = await process.stderr.read() - if err: - logger.error(err.decode().strip()) - - await process.wait() - logger.info(f"Command {command} {process}finished.") - - async def apply(self) -> []: - port = self.opts.get("port", 8718) - command = f"inspect eval {inspect_ai_task} --model openai/gpt-4 --model-base-url=http://0.0.0.0:{port}/proxy" - logger.info(f"Executing command: {command}") - - proc = asyncio.create_task(self._proc(command)) - is_empty = self.tools_inbox.empty() - await asyncio.sleep(2) - logger.info(f"Is inbox empty? {is_empty}") - while not self.tools_inbox.empty(): - ref = self.tools_inbox.get_nowait() - message, _, ready = ref["message"], ref["reply"], ref["ready"] - yield message - ready.set() - logger.info(f"{self.name} tool finished.") +import asyncio +import importlib.util +import os + +from agentic_security.logutils import logger + +inspect_ai_task = ( + __file__.replace("inspect_ai_tool.py", "inspect_ai_task.py") + .replace(os.getcwd(), "") + .strip("/") +) + + +class Module: + """:class:`Module` that runs Inspect AI evaluations against an LLM proxy. + + Launches the `Inspect AI `_ + evaluation framework via ``inspect eval``, targeting a local LLM proxy. + The module requires ``inspect_ai`` to be installed + (``pip install inspect_ai``). + + Configuration is passed via ``opts``: + + * ``port`` (int): LLM proxy base URL port. Defaults to ``8718``. + + Attributes: + tools_inbox: Async queue that receives evaluation results. + opts: Configuration dictionary. + """ + + name = "Inspect AI" + + def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue, opts: dict = {}): + self.tools_inbox = tools_inbox + if not self.is_tool_installed(): + logger.error( + "inspect_ai module is not installed. Please install it using 'pip install inspect_ai'" + ) + self.opts = opts + + def is_tool_installed(self) -> bool: + """Return ``True`` if the ``inspect_ai`` package is importable.""" + inspect_ai = importlib.util.find_spec("inspect_ai") + return inspect_ai is not None + + async def _proc(self, command): + env = os.environ.copy() + process = await asyncio.create_subprocess_shell( + command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=env, + shell=True, + ) + + logger.info(f"Started {command}") + + async for line in process.stdout: + logger.info(line.decode().strip()) + + err = await process.stderr.read() + if err: + logger.error(err.decode().strip()) + + await process.wait() + logger.info(f"Command {command} {process}finished.") + + async def apply(self) -> []: + port = self.opts.get("port", 8718) + command = f"inspect eval {inspect_ai_task} --model openai/gpt-4 --model-base-url=http://0.0.0.0:{port}/proxy" + logger.info(f"Executing command: {command}") + + proc = asyncio.create_task(self._proc(command)) + is_empty = self.tools_inbox.empty() + await asyncio.sleep(2) + logger.info(f"Is inbox empty? {is_empty}") + while not self.tools_inbox.empty(): + ref = self.tools_inbox.get_nowait() + message, _, ready = ref["message"], ref["reply"], ref["ready"] + yield message + ready.set() + logger.info(f"{self.name} tool finished.") await proc diff --git a/poetry.toml b/poetry.toml new file mode 100644 index 0000000..084377a --- /dev/null +++ b/poetry.toml @@ -0,0 +1,2 @@ +[virtualenvs] +create = false