From 58643f9d0a3c1ef9e78aac68a4c7ddd1162a3dce Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Mon, 29 Apr 2024 11:50:33 +0300
Subject: [PATCH] feat(Improve garak error handling):

---
 .../probe_data/modules/garak_tool.py          | 25 +++++++++++++++----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/agentic_security/probe_data/modules/garak_tool.py b/agentic_security/probe_data/modules/garak_tool.py
index a65560d..23864cf 100644
--- a/agentic_security/probe_data/modules/garak_tool.py
+++ b/agentic_security/probe_data/modules/garak_tool.py
@@ -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