Turn inherits from dict for serialisation

NVIDIA · Feb 17, 2025 · 957c1ec · 957c1ec
1 parent 49a9b17
commit 957c1ec
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 20 deletions.
diff --git a/garak/attempt.py b/garak/attempt.py
@@ -13,7 +13,7 @@
 roles = {"system", "user", "assistant"}
 
 
-class Turn:
+class Turn(dict):
     """Object to represent a single turn posed to or received from a generator
 
     Turns can be prompts, replies, system prompts. While many prompts are text,

diff --git a/garak/generators/nim.py b/garak/generators/nim.py
@@ -148,28 +148,31 @@ class Vision(NVOpenAIChat):
     modality = {"in": {"text", "image"}, "out": {"text"}}
 
     def _prepare_prompt(self, turn):
-        import base64
 
         text = turn.text
-        image_filename = turn.parts["image_filename"]
-        if image_filename is not None:
-            with open(image_filename, "rb") as f:
-                image_b64 = base64.b64encode(f.read()).decode()
-
-            if len(image_b64) > self.max_image_len:
-                logging.error(
-                    "Image %s exceeds length limit. To upload larger images, use the assets API (not yet supported)",
-                    image_filename,
+        if "image_filename" in turn.parts:
+            import base64
+
+            image_filename = turn.parts["image_filename"]
+            if image_filename is not None:
+                with open(image_filename, "rb") as f:
+                    image_b64 = base64.b64encode(f.read()).decode()
+
+                if len(image_b64) > self.max_image_len:
+                    logging.error(
+                        "Image %s exceeds length limit. To upload larger images, use the assets API (not yet supported)",
+                        image_filename,
+                    )
+                    return None
+
+                image_extension = turn.parts["image_filename"].split(".")[-1].lower()
+                if image_extension == "jpg":  # image/jpg is not a valid mimetype
+                    image_extension = "jpeg"
+                text = (
+                    text
+                    + f' <img src="data:image/{image_extension};base64,{image_b64}" />'
                 )
-                return None
-
-            image_extension = turn.parts["image_filename"].split(".")[-1].lower()
-            if image_extension == "jpg":  # image/jpg is not a valid mimetype
-                image_extension = "jpeg"
-            text = (
-                text + f' <img src="data:image/{image_extension};base64,{image_b64}" />'
-            )
-        turn.text = text
+            turn.text = text
         return turn