Skip to content

Commit

Permalink
updated tests to match new code (andrewyng#5)
Browse files Browse the repository at this point in the history 
* Create ruff.yml

CI check for linting

* updated ruff config

* updated tests
  • Loading branch information
@j-dominguez9
j-dominguez9 committed Jun 4, 2024
1 parent fd76cfb commit a8ece82
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 141 deletions.
14 changes: 5 additions & 9 deletions src/translation_agent/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 <SOURCE_TEXT> and </SOURCE_TEXT>, and output the translation.
Output only the translation and nothing else.
<SOURCE_TEXT>
{source_text}
</SOURCE_TEXT>
"""
{target_lang}:"""

prompt = translation_prompt.format(source_text=source_text)

Expand Down Expand Up @@ -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,
Expand Down
207 changes: 75 additions & 132 deletions tests/test_agent.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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():
Expand All @@ -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():
Expand All @@ -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(
Expand All @@ -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 <SOURCE_TEXT> and </SOURCE_TEXT> and output the translation.
Do not output anything other the translation.
{target_lang}:"""

<SOURCE_TEXT>
{source_text}
</SOURCE_TEXT>
"""
mock_get_completion_content.assert_called_once_with(
mock_get_completion.assert_called_once_with(
expected_prompt, system_message=expected_system_message
)

Expand All @@ -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 <SOURCE_TEXT> and <TRANSLATION>, are as follows:
The source text and initial translation, delimited by XML tags <SOURCE_TEXT></SOURCE_TEXT> and <TRANSLATION></TRANSLATION>, are as follows:
<SOURCE_TEXT>
{source_text}
</SOURCE_TEXT>
<TRANSLATION>
{translation1}
{translation_1}
</TRANSLATION>
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
)

Expand All @@ -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."
)

Expand All @@ -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 <SOURCE_TEXT>, <TRANSLATION> and <EXPERT_SUGGESTIONS> are as follows:
The source text, the initial translation, and the expert linguist suggestions are delimited by XML tags <SOURCE_TEXT></SOURCE_TEXT>, <TRANSLATION></TRANSLATION> and <EXPERT_SUGGESTIONS></EXPERT_SUGGESTIONS> \
as follows:
<SOURCE_TEXT>
This is a sample source text.
{example_data["source_text"]}
</SOURCE_TEXT>
<TRANSLATION>
Esta es una traducción de ejemplo.
{example_data["translation_1"]}
</TRANSLATION>
<EXPERT_SUGGESTIONS>
The translation is accurate but could be more fluent.
{example_data["reflection"]}
</EXPERT_SUGGESTIONS>
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
Expand All @@ -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
Expand Down Expand Up @@ -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()

0 comments on commit a8ece82

Please sign in to comment.