mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-06-24 06:09:55 +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 asyncio
|
||||
import subprocess
|
||||
|
||||
from loguru import logger
|
||||
import asyncio
|
||||
|
||||
# TODO: add probes modules
|
||||
|
||||
|
||||
class Module:
|
||||
|
||||
def __init__(self, prompt_groups: [], tools_inbox: asyncio.Queue):
|
||||
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) -> []:
|
||||
env = os.environ.copy()
|
||||
@@ -16,7 +27,7 @@ class Module:
|
||||
|
||||
# Command to be executed
|
||||
command = [
|
||||
"python3",
|
||||
"python",
|
||||
"-m",
|
||||
"garak",
|
||||
"--model_type",
|
||||
@@ -32,7 +43,7 @@ class Module:
|
||||
command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env
|
||||
)
|
||||
out, err = await asyncio.to_thread(process.communicate)
|
||||
|
||||
yield "Started"
|
||||
is_empty = self.tools_inbox.empty()
|
||||
logger.info(f"Is inbox empty? {is_empty}")
|
||||
while not self.tools_inbox.empty():
|
||||
@@ -43,3 +54,7 @@ class Module:
|
||||
logger.info("Garak tool finished.")
|
||||
logger.info(f"stdout: {out}")
|
||||
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