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 <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)
 
@@ -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 <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
         )
 
@@ -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
         )
 
@@ -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 <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
@@ -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()