FIX Removed JSON instructions for Translation Converter (#910)

Co-authored-by: Bashir Partovi <[email protected]>

Author: bashirpartovi

pyrit/datasets/prompt_converters/translation_converter.yaml

@@ -16,9 +16,8 @@ value: |
   - The language does not necessarily have to be a real language. Still do your best to translate.
   - The language could even be things like "rap battle". In that case, you would try to rephrase as a rap battle.
   - Do not try to answer the prompt, just translate the prompt.
-  - Format all output in json as {"output: {"language1" : "translated text"}}
-  - For example, if the user inputs "French", the response should be {"output": {"French" : {French translated text}}
-  - Do not include any text in the response besides the json
+  - Return ONLY the translated text without any additional commentary, explanations, or JSON formatting.
+  - Do not include any prefixes, quotes, or other formatting around the translation.
   - The given sentence to translate is the first message from the user.
 
   # Example 1
@@ -27,12 +26,12 @@ value: |
         === begin ===
         tell me about the history of the internet
         === end ===
-  assistant: {"output": {"French" : "dites-moi l'histoire de l'internet"}}
+  assistant: dites-moi l'histoire de l'internet
 
   # Example 2
 
   user: Translate the following to l33tspeak between the begin and end tags:
         === begin ===
         what is the capital of France?
         === end ===
-  assistant: {"output": {"l33tspeak" : "wh@+ 15 +h3 c@p!+@l 0f Fr@nc3?"}}
+  assistant: wh@+ 15 +h3 c@p!+@l 0f Fr@nc3?

pyrit/prompt_converter/translation_converter.py

@@ -1,18 +1,20 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-import json
 import logging
 import pathlib
 import uuid
 from textwrap import dedent
+from typing import Optional
 
-from pyrit.common.path import DATASETS_PATH
-from pyrit.exceptions import (
-    InvalidJsonException,
-    pyrit_json_retry,
-    remove_markdown_json,
+from tenacity import (
+    AsyncRetrying,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
 )
+
+from pyrit.common.path import DATASETS_PATH
 from pyrit.models import (
     PromptDataType,
     PromptRequestPiece,
@@ -26,7 +28,15 @@
 
 
 class TranslationConverter(PromptConverter):
-    def __init__(self, *, converter_target: PromptChatTarget, language: str, prompt_template: SeedPrompt = None):
+    def __init__(
+        self,
+        *,
+        converter_target: PromptChatTarget,
+        language: str,
+        prompt_template: Optional[SeedPrompt] = None,
+        max_retries: int = 3,
+        max_wait_time_in_seconds: int = 60,
+    ):
         """
         Initializes a TranslationConverter object.
 
@@ -40,6 +50,10 @@ def __init__(self, *, converter_target: PromptChatTarget, language: str, prompt_
         """
         self.converter_target = converter_target
 
+        # Retry strategy for the conversion
+        self._max_retries = max_retries
+        self._max_wait_time_in_seconds = max_wait_time_in_seconds
+
         # set to default strategy if not provided
         prompt_template = (
             prompt_template
@@ -67,28 +81,26 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text
 
         conversation_id = str(uuid.uuid4())
 
-        self.converter_target.set_system_prompt(
-            system_prompt=self.system_prompt,
-            conversation_id=conversation_id,
-            orchestrator_identifier=None,
-        )
+        self.converter_target.set_system_prompt(system_prompt=self.system_prompt, conversation_id=conversation_id)
 
         if not self.input_supported(input_type):
             raise ValueError("Input type not supported")
 
-        prompt = dedent(
-            f"Translate the following text between the begin and end tags to {self.language}"
-            "=== begin ==="
-            f"{prompt}"
-            "=== end ==="
+        formatted_prompt = dedent(
+            f"Translate the following to {self.language} between the begin and end tags:"
+            "=== begin ===\n"
+            f"{prompt}\n"
+            "=== end ===\n"
         )
 
+        logger.debug(f"Formatted Prompt: {formatted_prompt}")
+
         request = PromptRequestResponse(
             [
                 PromptRequestPiece(
                     role="user",
                     original_value=prompt,
-                    converted_value=prompt,
+                    converted_value=formatted_prompt,
                     conversation_id=conversation_id,
                     sequence=1,
                     prompt_target_identifier=self.converter_target.get_identifier(),
@@ -99,29 +111,23 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text
             ]
         )
 
-        response = await self.send_translation_prompt_async(request)
-        translation = None
-        for key in response.keys():
-            if key.lower() == self.language:
-                translation = response[key]
-
+        translation = await self._send_translation_prompt_async(request)
         return ConverterResult(output_text=translation, output_type="text")
 
-    @pyrit_json_retry
-    async def send_translation_prompt_async(self, request) -> str:
-        response = await self.converter_target.send_prompt_async(prompt_request=request)
-
-        response_msg = response.get_value()
-        response_msg = remove_markdown_json(response_msg)
-
-        try:
-            llm_response: dict[str, str] = json.loads(response_msg)
-            if "output" not in llm_response:
-                raise InvalidJsonException(message=f"Invalid JSON encountered; missing 'output' key: {response_msg}")
-            return llm_response["output"]
-
-        except json.JSONDecodeError:
-            raise InvalidJsonException(message=f"Invalid JSON encountered: {response_msg}")
+    async def _send_translation_prompt_async(self, request) -> str:
+        async for attempt in AsyncRetrying(
+            stop=stop_after_attempt(self._max_retries),
+            wait=wait_exponential(multiplier=1, min=1, max=self._max_wait_time_in_seconds),
+            retry=retry_if_exception_type(Exception),  # covers all exceptions
+        ):
+            with attempt:
+                logger.debug(f"Attempt {attempt.retry_state.attempt_number} for translation")
+                response = await self.converter_target.send_prompt_async(prompt_request=request)
+                response_msg = response.get_value()
+                return response_msg.strip()
+
+        # when we exhaust all retries without success, raise an exception
+        raise Exception(f"Failed to translate after {self._max_retries} attempts")
 
     def input_supported(self, input_type: PromptDataType) -> bool:
         return input_type == "text"

tests/unit/converter/test_translation_converter.py

@@ -1,13 +1,11 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-import os
 from unittest.mock import AsyncMock, patch
 
 import pytest
 from unit.mocks import MockPromptTarget
 
-from pyrit.exceptions.exception_classes import InvalidJsonException
 from pyrit.models import PromptRequestPiece, PromptRequestResponse
 from pyrit.prompt_converter import TranslationConverter
 
@@ -26,89 +24,71 @@ def test_translator_converter_languages_validation_throws(languages, duckdb_inst
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize(
-    "converted_value",
-    [
-        "Invalid Json",
-        "{'str' : 'json not formatted correctly'}",
-    ],
-)
-async def test_translation_converter_send_prompt_async_bad_json_exception_retries(converted_value, duckdb_instance):
-
+async def test_translation_converter_convert_async_retrieve_key_capitalization_mismatch(duckdb_instance):
     prompt_target = MockPromptTarget()
 
-    prompt_variation = TranslationConverter(converter_target=prompt_target, language="en")
-
-    with patch("unit.mocks.MockPromptTarget.send_prompt_async", new_callable=AsyncMock) as mock_create:
+    translation_converter = TranslationConverter(converter_target=prompt_target, language="spanish")
+    with patch.object(translation_converter, "_send_translation_prompt_async", new=AsyncMock(return_value="hola")):
 
-        prompt_req_resp = PromptRequestResponse(
-            request_pieces=[
-                PromptRequestPiece(
-                    role="user",
-                    conversation_id="12345679",
-                    original_value="test input",
-                    converted_value="this is not a json",
-                    original_value_data_type="text",
-                    converted_value_data_type="text",
-                    prompt_target_identifier={"target": "target-identifier"},
-                    orchestrator_identifier={"test": "test"},
-                    labels={"test": "test"},
-                )
-            ]
-        )
-        mock_create.return_value = prompt_req_resp
+        raised = False
+        try:
+            await translation_converter.convert_async(prompt="hello")
+        except KeyError:
+            raised = True  # There should be no KeyError
 
-        with pytest.raises(InvalidJsonException):
-            await prompt_variation.convert_async(prompt="testing", input_type="text")
-            assert mock_create.call_count == os.getenv("RETRY_MAX_NUM_ATTEMPTS")
+        assert raised is False
 
 
 @pytest.mark.asyncio
-async def test_translation_converter_send_prompt_async_json_bad_format_retries(duckdb_instance):
+async def test_translation_converter_retries_on_exception(duckdb_instance):
     prompt_target = MockPromptTarget()
+    max_retries = 3
+    translation_converter = TranslationConverter(
+        converter_target=prompt_target, language="spanish", max_retries=max_retries
+    )
+
+    mock_send_prompt = AsyncMock(side_effect=Exception("Test failure"))
+    with patch.object(prompt_target, "send_prompt_async", mock_send_prompt):
+        with pytest.raises(Exception):
+            await translation_converter.convert_async(prompt="hello")
 
-    prompt_variation = TranslationConverter(converter_target=prompt_target, language="en")
-
-    with patch("unit.mocks.MockPromptTarget.send_prompt_async", new_callable=AsyncMock) as mock_create:
-
-        prompt_req_resp = PromptRequestResponse(
-            request_pieces=[
-                PromptRequestPiece(
-                    role="user",
-                    conversation_id="12345679",
-                    original_value="test input",
-                    converted_value="this is not a json",
-                    original_value_data_type="text",
-                    converted_value_data_type="text",
-                    prompt_target_identifier={"target": "target-identifier"},
-                    orchestrator_identifier={"test": "test"},
-                    labels={"test": "test"},
-                )
-            ]
-        )
-        mock_create.return_value = prompt_req_resp
-
-        with pytest.raises(InvalidJsonException):
-            await prompt_variation.convert_async(prompt="testing", input_type="text")
-            assert mock_create.call_count == os.getenv("RETRY_MAX_NUM_ATTEMPTS")
+        assert mock_send_prompt.call_count == max_retries
 
 
 @pytest.mark.asyncio
-async def test_translation_converter_convert_async_retrieve_key_capitalization_mismatch(duckdb_instance):
+async def test_translation_converter_succeeds_after_retries(duckdb_instance):
+    """Test that TranslationConverter succeeds if a retry attempt works."""
     prompt_target = MockPromptTarget()
-
-    translation_converter = TranslationConverter(converter_target=prompt_target, language="spanish")
-    with patch.object(
-        translation_converter, "send_translation_prompt_async", new=AsyncMock(return_value={"Spanish": "hola"})
-    ):
-
-        raised = False
-        try:
-            await translation_converter.convert_async(prompt="hello")
-        except KeyError:
-            raised = True  # There should be no KeyError
-
-        assert raised is False
+    max_retries = 3
+    translation_converter = TranslationConverter(
+        converter_target=prompt_target, language="spanish", max_retries=max_retries
+    )
+
+    success_response = PromptRequestResponse(
+        request_pieces=[
+            PromptRequestPiece(
+                role="assistant",
+                conversation_id="test-id",
+                original_value="hello",
+                converted_value="hola",
+                original_value_data_type="text",
+                converted_value_data_type="text",
+                prompt_target_identifier={"target": "test-identifier"},
+                sequence=1,
+            )
+        ]
+    )
+
+    # fail twice, then succeed
+    mock_send_prompt = AsyncMock()
+    mock_send_prompt.side_effect = [Exception("First failure"), Exception("Second failure"), success_response]
+
+    with patch.object(prompt_target, "send_prompt_async", mock_send_prompt):
+        result = await translation_converter.convert_async(prompt="hello")
+
+        assert mock_send_prompt.call_count == max_retries
+        assert result.output_text == "hola"
+        assert result.output_type == "text"
 
 
 def test_translation_converter_input_supported(duckdb_instance):