From a8ece8226bac80109d167fe87a69854c9bf7443c Mon Sep 17 00:00:00 2001
From: Joaquin Dominguez <83036592+j-dominguez9@users.noreply.github.com>
Date: Tue, 4 Jun 2024 16:31:21 -0400
Subject: [PATCH] updated tests to match new code (#5)
* Create ruff.yml
CI check for linting
* updated ruff config
* updated tests
---
src/translation_agent/utils.py | 14 +--
tests/test_agent.py | 207 ++++++++++++---------------------
2 files changed, 80 insertions(+), 141 deletions(-)
diff --git a/src/translation_agent/utils.py b/src/translation_agent/utils.py
index 42fcbf7..0583821 100755
--- a/src/translation_agent/utils.py
+++ b/src/translation_agent/utils.py
@@ -87,15 +87,11 @@ def one_chunk_initial_translation(
system_message = f"You are an expert linguist, specializing in translation from {source_lang} to {target_lang}."
- translation_prompt = f"""Your task is to provide a professional translation of a text from {source_lang} to {target_lang}.
+ translation_prompt = f"""This is an {source_lang} to {target_lang} translation, please provide the {target_lang} translation for this text. \
+Do not provide any explanations or text apart from the translation.
+{source_lang}: {source_text}
- Translate the text below, delimited by XML tags and , and output the translation.
- Output only the translation and nothing else.
-
-
- {source_text}
-
- """
+{target_lang}:"""
prompt = translation_prompt.format(source_text=source_text)
@@ -148,7 +144,7 @@ def one_chunk_reflect_on_translation(
Write a list of specific, helpful and constructive suggestions for improving the translation.
Each suggestion should address one specific part of the translation.
-Output only the suggestions and nothing else."""
+Output only the suggestions and nothing else."""
prompt = reflection_prompt.format(
source_lang=source_lang,
diff --git a/tests/test_agent.py b/tests/test_agent.py
index 5188ec4..88461b1 100644
--- a/tests/test_agent.py
+++ b/tests/test_agent.py
@@ -4,15 +4,15 @@
import openai
import pytest
-import spacy
from dotenv import load_dotenv
-from translation_agent.util import find_sentence_starts
-from translation_agent.util import get_completion
-from translation_agent.util import num_tokens_in_string
-from translation_agent.util import one_chunk_improve_translation
-from translation_agent.util import one_chunk_initial_translation
-from translation_agent.util import one_chunk_reflect_on_translation
-from translation_agent.util import one_chunk_translate_text
+
+# from translation_agent.utils import find_sentence_starts
+from translation_agent.utils import get_completion
+from translation_agent.utils import num_tokens_in_string
+from translation_agent.utils import one_chunk_improve_translation
+from translation_agent.utils import one_chunk_initial_translation
+from translation_agent.utils import one_chunk_reflect_on_translation
+from translation_agent.utils import one_chunk_translate_text
load_dotenv()
@@ -36,17 +36,8 @@ def test_get_completion_json_mode_api_call():
# Assert that the result is not None
assert result is not None
- # Assert that the result has the expected keys
- assert "id" in dir(result)
- assert "object" in dir(result)
- assert "created" in dir(result)
- assert "model" in dir(result)
- assert "choices" in dir(result)
-
- # Assert that the result has the expected response format
- assert result.object == "chat.completion"
- assert result.choices[0].message.role == "assistant"
- assert isinstance(json.loads(result.choices[0].message.content), dict)
+ # Assert that it can be transformed to dictionary (json)
+ assert isinstance(json.loads(result), dict)
def test_get_completion_non_json_mode_api_call():
@@ -65,16 +56,8 @@ def test_get_completion_non_json_mode_api_call():
# Assert that the result is not None
assert result is not None
- # Assert that the result has the expected keys
- assert "id" in dir(result)
- assert "created" in dir(result)
- assert "model" in dir(result)
- assert "choices" in dir(result)
-
# Assert that the result has the expected response format
- assert result.object == "chat.completion"
- assert result.choices[0].message.role == "assistant"
- assert isinstance(result.choices[0].message.content, str)
+ assert isinstance(result, str)
def test_one_chunk_initial_translation():
@@ -86,9 +69,9 @@ def test_one_chunk_initial_translation():
# Mock the get_completion_content function
with patch(
- "translation_agent.util.get_completion_content"
- ) as mock_get_completion_content:
- mock_get_completion_content.return_value = expected_translation
+ "translation_agent.utils.get_completion"
+ ) as mock_get_completion:
+ mock_get_completion.return_value = expected_translation
# Call the function with test data
translation = one_chunk_initial_translation(
@@ -99,17 +82,14 @@ def test_one_chunk_initial_translation():
assert translation == expected_translation
# Assert the get_completion_content function was called with the correct arguments
- expected_system_message = f"You are an expert language translator, specializing in {source_lang} to {target_lang} translation."
- expected_prompt = f"""You task is provide a professional translation of a text from {source_lang} to {target_lang}.
+ expected_system_message = f"You are an expert linguist, specializing in translation from {source_lang} to {target_lang}."
+ expected_prompt = f"""This is an {source_lang} to {target_lang} translation, please provide the {target_lang} translation for this text. \
+Do not provide any explanations or text apart from the translation.
+{source_lang}: {source_text}
-Translate the text below, delimited by XML tags and and output the translation.
-Do not output anything other the translation.
+{target_lang}:"""
-
-{source_text}
-
-"""
- mock_get_completion_content.assert_called_once_with(
+ mock_get_completion.assert_called_once_with(
expected_prompt, system_message=expected_system_message
)
@@ -118,50 +98,53 @@ def test_one_chunk_reflect_on_translation():
# Define test data
source_lang = "English"
target_lang = "Spanish"
+ country = "Mexico"
source_text = "This is a sample source text."
- translation1 = "Este es un texto de origen de muestra."
+ translation_1 = "Este es un texto de origen de muestra."
# Define the expected reflection
expected_reflection = "The translation is accurate and conveys the meaning of the source text well. However, here are a few suggestions for improvement:\n\n1. Consider using 'texto fuente' instead of 'texto de origen' for a more natural translation of 'source text'.\n2. Add a definite article before 'texto fuente' to improve fluency: 'Este es un texto fuente de muestra.'\n3. If the context allows, you could also use 'texto de ejemplo' as an alternative translation for 'sample text'."
# Mock the get_completion_content function
with patch(
- "translation_agent.util.get_completion_content"
- ) as mock_get_completion_content:
- mock_get_completion_content.return_value = expected_reflection
+ "translation_agent.utils.get_completion"
+ ) as mock_get_completion:
+ mock_get_completion.return_value = expected_reflection
# Call the function with test data
reflection = one_chunk_reflect_on_translation(
- source_lang, target_lang, source_text, translation1
+ source_lang, target_lang, source_text, translation_1, country
)
# Assert that the reflection matches the expected reflection
assert reflection == expected_reflection
# Assert that the get_completion_content function was called with the correct arguments
- expected_prompt = f"""Your task is to carefully read a source text and a translation from {source_lang} to {target_lang}, and then give constructive criticism and helpful suggestions for the translation.
+ expected_prompt = f"""Your task is to carefully read a source text and a translation from {source_lang} to {target_lang}, and then give constructive criticism and helpful suggestions to improve the translation. \
+The final style and tone of the translation should match the style of {target_lang} colloquially spoken in {country}.
-The source text and initial translation, delimited by XML tags and , are as follows:
+The source text and initial translation, delimited by XML tags and , are as follows:
{source_text}
-{translation1}
+{translation_1}
-When writing suggestions, pay attention to whether there are ways to improve the translation's \
-(i) accuracy (by correcting errors of addition, mistranslation, omission, untranslated text),
-(ii) fluency (grammar, inconsistency, punctuation, register, spelling), \
-(iii) style (fix awkward wording),
-(iv) terminology (inappropriate for context, inconsistent use), or \
-(v) other errors.
+When writing suggestions, pay attention to whether there are ways to improve the translation's \n\
+(i) accuracy (by correcting errors of addition, mistranslation, omission, or untranslated text),\n\
+(ii) fluency (by applying {target_lang} grammar, spelling and punctuation rules, and ensuring there are no unnecessary repetitions),\n\
+(iii) style (by ensuring the translations reflect the style of the source text and takes into account any cultural context),\n\
+(iv) terminology (by ensuring terminology use is consistent and reflects the source text domain; and by only ensuring you use equivalent idioms {target_lang}).\n\
Write a list of specific, helpful and constructive suggestions for improving the translation.
-Each suggestion should address one specific part of the translation."""
- expected_system_message = f"You are an expert language translator and mentor, specializing in {source_lang} to {target_lang} translation."
- mock_get_completion_content.assert_called_once_with(
+Each suggestion should address one specific part of the translation.
+Output only the suggestions and nothing else."""
+ expected_system_message = f"You are an expert linguist specializing in translation from {source_lang} to {target_lang}. \
+You will be provided with a source text and its translation and your goal is to improve the translation."
+ mock_get_completion.assert_called_once_with(
expected_prompt, system_message=expected_system_message
)
@@ -172,17 +155,15 @@ def example_data():
"source_lang": "English",
"target_lang": "Spanish",
"source_text": "This is a sample source text.",
- "translation1": "Esta es una traducción de ejemplo.",
+ "translation_1": "Esta es una traducción de ejemplo.",
"reflection": "The translation is accurate but could be more fluent.",
}
-@patch("translation_agent.util.get_completion_content")
-def test_one_chunk_improve_translation(
- mock_get_completion_content, example_data
-):
+@patch("translation_agent.utils.get_completion")
+def test_one_chunk_improve_translation(mock_get_completion, example_data):
# Set up the mock return value for get_completion_content
- mock_get_completion_content.return_value = (
+ mock_get_completion.return_value = (
"Esta es una traducción de ejemplo mejorada."
)
@@ -191,73 +172,77 @@ def test_one_chunk_improve_translation(
example_data["source_lang"],
example_data["target_lang"],
example_data["source_text"],
- example_data["translation1"],
+ example_data["translation_1"],
example_data["reflection"],
)
# Assert that the function returns the expected translation
assert result == "Esta es una traducción de ejemplo mejorada."
- # Assert that get_completion_content was called with the expected arguments
- expected_prompt = """Your task is to carefully read, then improve, a translation from English to Spanish, taking into
-account a set of expert suggestions and constructive critisms.
+ # Assert that get_completion was called with the expected arguments
+ expected_prompt = f"""Your task is to carefully read, then edit, a translation from {example_data["source_lang"]} to {example_data["target_lang"]}, taking into
+account a list of expert suggestions and constructive criticisms.
-The source text, initial translation, and expert suggestions, delimited by XML tags , and are as follows:
+The source text, the initial translation, and the expert linguist suggestions are delimited by XML tags , and \
+as follows:
-This is a sample source text.
+{example_data["source_text"]}
-Esta es una traducción de ejemplo.
+{example_data["translation_1"]}
-The translation is accurate but could be more fluent.
+{example_data["reflection"]}
-Taking into account the expert suggestions rewrite the translation to improve it, paying attention
-to whether there are ways to improve the translation's \
-(i) accuracy (by correcting errors of addition, mistranslation, omission, untranslated text),
-(ii) fluency (grammar, inconsistency, punctuation, register, spelling), \
-(iii) style (fix awkward wording),
+Please take into account the expert suggestions when editing the translation. Edit the translation by ensuring:
+
+(i) accuracy (by correcting errors of addition, mistranslation, omission, or untranslated text),
+(ii) fluency (by applying Spanish grammar, spelling and punctuation rules and ensuring there are no unnecessary repetitions), \
+(iii) style (by ensuring the translations reflect the style of the source text)
(iv) terminology (inappropriate for context, inconsistent use), or \
-(v) other errors. Output the list of suggestions in JSON, using the key "suggestions".
+(v) other errors.
-Output the new translation, and nothing else."""
+Output only the new translation and nothing else."""
- expected_system_message = "You are an expert language translator, specializing in English to Spanish translation."
+ expected_system_message = f"You are an expert linguist, specializing in translation editing from English to Spanish."
- mock_get_completion_content.assert_called_once_with(
+ mock_get_completion.assert_called_once_with(
expected_prompt, expected_system_message
)
def test_one_chunk_translate_text(mocker):
# Define test data
- source_lang = "en"
- target_lang = "es"
+ source_lang = "English"
+ target_lang = "Spanish"
+ country = "Mexico"
source_text = "Hello, how are you?"
- translation1 = "Hola, ¿cómo estás?"
+ translation_1 = "Hola, ¿cómo estás?"
reflection = "The translation looks good, but it could be more formal."
translation2 = "Hola, ¿cómo está usted?"
# Mock the helper functions
mock_initial_translation = mocker.patch(
- "translation_agent.util.one_chunk_initial_translation",
- return_value=translation1,
+ "translation_agent.utils.one_chunk_initial_translation",
+ return_value=translation_1,
)
mock_reflect_on_translation = mocker.patch(
- "translation_agent.util.one_chunk_reflect_on_translation",
+ "translation_agent.utils.one_chunk_reflect_on_translation",
return_value=reflection,
)
mock_improve_translation = mocker.patch(
- "translation_agent.util.one_chunk_improve_translation",
+ "translation_agent.utils.one_chunk_improve_translation",
return_value=translation2,
)
# Call the function being tested
- result = one_chunk_translate_text(source_lang, target_lang, source_text)
+ result = one_chunk_translate_text(
+ source_lang, target_lang, source_text, country
+ )
# Assert the expected result
assert result == translation2
@@ -267,45 +252,13 @@ def test_one_chunk_translate_text(mocker):
source_lang, target_lang, source_text
)
mock_reflect_on_translation.assert_called_once_with(
- source_lang, target_lang, source_text, translation1
+ source_lang, target_lang, source_text, translation_1, country
)
mock_improve_translation.assert_called_once_with(
- source_lang, target_lang, source_text, translation1, reflection
+ source_lang, target_lang, source_text, translation_1, reflection
)
-@pytest.fixture(scope="module")
-def english_model():
- return spacy.load("en_core_web_sm")
-
-
-def test_find_sentence_starts(english_model):
- # Test case 1: Single sentence
- text1 = "This is a single sentence."
- expected_output1 = [0]
- assert find_sentence_starts(text1) == expected_output1
-
- # Test case 2: Multiple sentences
- text2 = "This is the first sentence. This is the second sentence. And this is the third sentence."
- expected_output2 = [0, 28, 57]
- assert find_sentence_starts(text2) == expected_output2
-
- # Test case 3: Empty string
- text3 = ""
- expected_output3 = []
- assert find_sentence_starts(text3) == expected_output3
-
- # Test case 4: Text with punctuation
- text4 = "Hello, how are you? I'm doing fine! Thanks for asking."
- expected_output4 = [0, 20, 36]
- assert find_sentence_starts(text4) == expected_output4
-
- # Test case 5: Text with multiple spaces and newlines
- text5 = " This is a sentence. \n\nAnother sentence here. \n One more sentence. "
- expected_output5 = [0, 26, 51]
- assert find_sentence_starts(text5) == expected_output5
-
-
def test_num_tokens_in_string():
# Test case 1: Empty string
assert num_tokens_in_string("") == 0
@@ -334,13 +287,3 @@ def test_num_tokens_in_string():
assert (
num_tokens_in_string("Hello, world!", encoding_name="p50k_base") == 4
)
-
-
-# def test_invalid_encoding():
-# # Test case 7: Invalid encoding name
-# with pytest.raises(KeyError):
-# num_tokens_in_string("Hello, world!", encoding_name="invalid_encoding")
-
-
-# result = get_completion(prompt="this is a test, reply with success")
-# breakpoint()