mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-08-09 03:06:05 +02:00
feat(Improve garak error handling):
This commit is contained in:
@@ -1,14 +1,25 @@
|
|||||||
import subprocess
|
import asyncio
|
||||||
|
import importlib.util
|
||||||
import os
|
import os
|
||||||
import asyncio
|
import subprocess
|
||||||
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
import asyncio
|
|
||||||
|
# TODO: add probes modules
|
||||||
|
|
||||||
|
|
||||||
class Module:
|
class Module:
|
||||||
|
|
||||||
def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue):
|
def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue):
|
||||||
self.tools_inbox = tools_inbox
|
self.tools_inbox = tools_inbox
|
||||||
|
if not self.is_garak_installed():
|
||||||
|
logger.error(
|
||||||
|
"Garak module is not installed. Please install it using 'pip install garak'"
|
||||||
|
)
|
||||||
|
|
||||||
|
def is_garak_installed(self) -> bool:
|
||||||
|
garak_spec = importlib.util.find_spec("garak")
|
||||||
|
return garak_spec is not None
|
||||||
|
|
||||||
async def apply(self) -> []:
|
async def apply(self) -> []:
|
||||||
env = os.environ.copy()
|
env = os.environ.copy()
|
||||||
@@ -16,7 +27,7 @@ class Module:
|
|||||||
|
|
||||||
# Command to be executed
|
# Command to be executed
|
||||||
command = [
|
command = [
|
||||||
"python3",
|
"python",
|
||||||
"-m",
|
"-m",
|
||||||
"garak",
|
"garak",
|
||||||
"--model_type",
|
"--model_type",
|
||||||
@@ -32,7 +43,7 @@ class Module:
|
|||||||
command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env
|
command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env
|
||||||
)
|
)
|
||||||
out, err = await asyncio.to_thread(process.communicate)
|
out, err = await asyncio.to_thread(process.communicate)
|
||||||
|
yield "Started"
|
||||||
is_empty = self.tools_inbox.empty()
|
is_empty = self.tools_inbox.empty()
|
||||||
logger.info(f"Is inbox empty? {is_empty}")
|
logger.info(f"Is inbox empty? {is_empty}")
|
||||||
while not self.tools_inbox.empty():
|
while not self.tools_inbox.empty():
|
||||||
@@ -43,3 +54,7 @@ class Module:
|
|||||||
logger.info("Garak tool finished.")
|
logger.info("Garak tool finished.")
|
||||||
logger.info(f"stdout: {out}")
|
logger.info(f"stdout: {out}")
|
||||||
logger.error(f"exit code: {process.returncode}")
|
logger.error(f"exit code: {process.returncode}")
|
||||||
|
if process.returncode != 0:
|
||||||
|
logger.error(f"Error executing command: {command}")
|
||||||
|
logger.error(f"err: {err}")
|
||||||
|
return
|
||||||
|
|||||||
Reference in New Issue
Block a user