feat(Improve garak error handling):

This commit is contained in:
Alexander Myasoedov
2024-04-29 11:50:33 +03:00
parent 1fbcde8bbb
commit 58643f9d0a
@@ -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