initial delivery

.env.sample

@@ -0,0 +1,3 @@
+OPENAI_API_KEY="sk-xxxxx"    # replace "sk-xxxxx" with your secret OpenAI API key 
+OPENAI_BASE_URL=""           # replace "" with your API base URL, eg. "https://api.openai.com"
+OPENAI_MODEL = ""            # replace "" with your model name, eg. deepseek-chat

.gitignore

@@ -0,0 +1,8 @@
+cache_dir
+.env
+.venv
+__pycache__
+poetry.lock
+floresp-v2.0-rc.3
+*cache
+wmt

.pre-commit-config.yaml

@@ -0,0 +1,23 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.5.0
+    hooks:
+      - id: trailing-whitespace
+        exclude: tests
+      - id: end-of-file-fixer
+      - id: check-merge-conflict
+      - id: check-case-conflict
+      - id: check-json
+      - id: check-toml
+        exclude: tests/fixtures/invalid_lock/poetry\.lock
+      - id: check-yaml
+      - id: pretty-format-json
+        args: [--autofix, --no-ensure-ascii, --no-sort-keys]
+      - id: check-ast
+      - id: debug-statements
+      - id: check-docstring-first
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.3.5
+    hooks:
+      - id: ruff
+      - id: ruff-format

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 pisces76
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,13 +1,51 @@
----
-title: Translation Agent UI
-emoji: ⚡
-colorFrom: gray
-colorTo: pink
-sdk: gradio
-sdk_version: 4.36.1
-app_file: app.py
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Translation Agent UI: A UI for Andrew's Translation Agent
+
+This is a gradio UI for Andrew Ng's [Translation Agent](https://github.com/andrewyng/translation-agent). The main features are:
+1. Look inside the reflection agentic workflow during tranlation; 
+2. Support other LLM models as translation engine; 
+3. Highlight the differences between initial translation and final translation; 
+
+<img width="1280" src="./examples/demo-zh_en.png" alt="Preview"/>
+<img width="1280" src="./examples/demo-en_zh.png" alt="Preview"/>
+
+## Getting Started
+
+To get started with `translation-agent-UI`, follow these steps:
+
+### Installation:
+- The Poetry package manager is required for installation. [Poetry Installation](https://python-poetry.org/docs/#installation) Depending on your environment, this might work:
+
+```bash
+pip install poetry 
+```
+
+- A .env file with a OPENAI_API_KEY is required to run the workflow. Please copy .env.sample to .env and modified the parameters. Following is an example to use deekseek LLM as a translation engine.
+OPENAI_API_KEY="sk-xxxx"
+OPENAI_BASE_URL="https://api.deepseek.com"
+OPENAI_MODEL = "deepseek-chat"
+
+- Clone the repo and install the dependencies:
+```bash
+git clone https://github.com/andrewyng/translation-agent.git
+cd translation-agent
+poetry install
+poetry shell # activates virtual environment
+```
+### Usage:
+python gr_app.py  
+The web server will be running on local URL:  http://127.0.0.1:7860 by default.
+
+you can also use CLI to 
+```python
+import translation_agent as ta
+source_lang, target_lang, country = "English", "Spanish", "Mexico"
+translation = ta.translate(source_lang, target_lang, source_text, country)
+```
+See examples/example_script.py for an example script to try out.
+
+## License
+
+Translation Agent UI is released under the **MIT License**. You are free to use, modify, and distribute the code
+for both commercial and non-commercial purposes.
+
+

examples/README.md

@@ -0,0 +1,16 @@
+# Examples
+
+This directory contains example scripts demonstrating the usage of the `translation-agent` workflow.
+
+## Contents
+- `example_script.py`: A simple script showing how to perform machine translation using the package.
+- `sample-texts/`: A directory containing a few sample texts from The Batch letters written by Andrew and Data Points summaries found on the [DeepLearning.ai website](https://www.deeplearning.ai/the-batch/tag/data-points/).
+
+## Usage
+To run the example scripts, ensure that you have installed the `translation-agent` package and have activated your virtual environment. Then run:
+
+```python
+python example_script.py
+```
+
+If you have any questions or encounter any issues, please feel free to open an issue on the Github repository.

examples/example_script.py

@@ -0,0 +1,26 @@
+import os
+
+import translation_agent as ta
+
+
+if __name__ == "__main__":
+    source_lang, target_lang, country = "English", "Spanish", "Mexico"
+
+    relative_path = "sample-texts/sample-short1.txt"
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+
+    full_path = os.path.join(script_dir, relative_path)
+
+    with open(full_path, encoding="utf-8") as file:
+        source_text = file.read()
+
+    print(f"Source text:\n\n{source_text}\n------------\n")
+
+    translation = ta.translate(
+        source_lang=source_lang,
+        target_lang=target_lang,
+        source_text=source_text,
+        country=country,
+    )
+
+    print(f"Translation:\n\n{translation}")

examples/sample-texts/data_points_samples.json

@@ -0,0 +1 @@
+[{"text": "Paid ChatGPT users can now upload files directly from Google Drive and Microsoft OneDrive, interact with tables and charts using natural language, and customize charts for presentations. When users upload or import a data file, ChatGPT can now write and execute Python code to analyze or visualize that data on users\u2019 behalf. These features may make it easier for those with limited coding skills to conduct in-depth analyses and let experts save time on routine data tasks."}, {"text": "Reddit\u2019s vast forums will be used to power ChatGPT and other AI products. The collaboration will give Reddit new AI-powered features for its users and moderators, while OpenAI will advertise on Reddit. (Full terms were undisclosed.) OpenAI now has deals with global newspapers, software forums, and a wide variety of other publishers, giving it special access to timely and high-quality training material."}, {"text": "ZeroGPU is accessible through Hugging Face\u2019s Spaces platform, which already hosts over 300,000 AI demos. The shared Nvidia A100s can be used concurrently by multiple users or applications; unutilized capacity will be made available to others. HuggingFace\u2019s goal is to counter tech giants and closed models\u2019 centralization by making state-of-the-art AI technologies more accessible."}, {"text": "Chameleon can natively process both text and images together, allowing it to perform a wide range of mixed-modal tasks with impressive results. Meta\u2019s researchers say the key is Chameleon\u2019s fully token-based architecture (representing images as well as texts as tokens) and training on datasets that combine text with images. Chameleon outperforms many leading and specialized models (including GPT-4V and Gemini Pro) when answering questions about images, describing pictures, writing relevant text, and creating images from text prompts.\u00a0"}, {"text": "Google\u2019s AI-assisted, browser-based integrated development environment (IDE) offers now-familiar features like code completion, debugging tools, and a chat-assisted sidebar, all powered by Gemini. Whenever IDX modifies snippets or suggests new code, it also links back to the original source and its associated license, ensuring proper attribution. Although Google is entering a competitive market, IDX aims to attract developers by showcasing Gemini\u2019s AI advancements and integrating with the company\u2019s cloud services."}, {"text": "The tool aims to solve new users\u2019 \u201cblank page problem\u201d by providing a starting point for testing and iteration, incorporating best practices like chain of thought and separating data from instructions. Users can access the prompt generator directly on the Console or analyze the underlying prompt and architecture using a Google Colab notebook. The generator addresses a common challenge for AI users: efficiently crafting effective (and often larger and more complex) prompts that yield high-quality results."}, {"text": "ElevenLabs Reader: AI Audio is the billion-dollar AI voice cloning startup\u2019s first consumer app. The free app can read web pages, PDFs, and other documents aloud using a selection of 11 AI-generated voices. The app marks ElevenLabs\u2019 expansion into the broader AI voice market beyond its current focus on entertainment and media production."}, {"text": "Microsoft reportedly asked hundreds of its China-based employees working on cloud computing and AI to consider relocating to other countries. One source said Microsoft offered 700 to 800 Chinese engineers the opportunity to transfer to the U.S., Ireland, Australia, or New Zealand. The move comes as the U.S. government tightens restrictions on China\u2019s access to advanced technology, citing concerns over potential military applications and cybersecurity threats."}, {"text": "Abu Dhabi\u2019s Technology Innovation Institute released Falcon 2, a family of large language models that includes Falcon 2 11B and Falcon 2 11B VLM. The latter is the institute\u2019s first multimodal model, capable of converting visual inputs into textual outputs. Both models are Apache 2.0 open-source, multilingual, and perform on par with Gemma 7B and better than Llama 3 8B according to benchmarks and HuggingFace leaderboards."}]

examples/sample-texts/sample-long1.txt

@@ -0,0 +1,16 @@
+Last week, I spoke about AI and regulation at the U.S. Capitol at an event that was attended by legislative and business leaders. I’m encouraged by the progress the open source community has made fending off regulations that would have stifled innovation. But opponents of open source are continuing to shift their arguments, with the latest worries centering on open source's impact on national security. I hope we’ll all keep protecting open source!
+
+Based on my conversations with legislators, I’m encouraged by the progress the U.S. federal government has made getting a realistic grasp of AI’s risks. To be clear, guardrails are needed. But they should be applied to AI applications, not to general-purpose AI technology.
+
+Nonetheless, as I wrote previously, some companies are eager to limit open source, possibly to protect the value of massive investments they’ve made in proprietary models and to deter competitors. It has been fascinating to watch their arguments change over time.
+
+For instance, about 12 months ago, the Center For AI Safety’s “Statement on AI Risk” warned that AI could cause human extinction and stoked fears of AI taking over. This alarmed leaders in Washington. But many people in AI pointed out that this dystopian science-fiction scenario has little basis in reality. About six months later, when I testified at the U.S. Senate’s AI Insight forum, legislators no longer worried much about an AI takeover.
+
+Then the opponents of open source shifted gears. Their leading argument shifted to the risk of AI helping to create bioweapons. Soon afterward, OpenAI and RAND showed that current AI does not significantly increase the ability of malefactors to build bioweapons. This fear of AI-enabled bioweapons has diminished. To be sure, the possibility that bad actors could use bioweapons — with or without AI — remains a topic of great international concern.
+
+
+The latest argument for blocking open source AI has shifted to national security. AI is useful for both economic competition and warfare, and open source opponents say the U.S. should make sure its adversaries don’t have access to the latest foundation models. While I don’t want authoritarian governments to use AI, particularly to wage unjust wars, the LLM cat is out of the bag, and authoritarian countries will fill the vacuum if democratic nations limit access. When, some day, a child somewhere asks an AI system questions about democracy, the role of a free press, or the function of an independent judiciary in preserving the rule of law, I would like the AI to reflect democratic values rather than favor authoritarian leaders’ goals over, say, human rights.
+
+I came away from Washington optimistic about the progress we’ve made. A  year ago, legislators seemed to me to spend 80% of their time talking about guardrails for AI and 20% about investing in innovation. I was delighted that the ratio has flipped, and there was far more talk of investing in innovation.
+
+Looking beyond the U.S. federal government, there are many jurisdictions globally. Unfortunately, arguments in favor of  regulations that would stifle AI development continue to proliferate. But I’ve learned from my trips to Washington and other nations’ capitals that talking to regulators does have an impact. If you get a chance to talk to a regulator at any level, I hope you’ll do what you can to help governments better understand AI.

examples/sample-texts/sample-short1.txt

@@ -0,0 +1,3 @@
+Last week, I spoke about AI and regulation at the U.S. Capitol at an event that was attended by legislative and business leaders. I’m encouraged by the progress the open source community has made fending off regulations that would have stifled innovation. But opponents of open source are continuing to shift their arguments, with the latest worries centering on open source's impact on national security. I hope we’ll all keep protecting open source!
+
+Based on my conversations with legislators, I’m encouraged by the progress the U.S. federal government has made getting a realistic grasp of AI’s risks. To be clear, guardrails are needed. But they should be applied to AI applications, not to general-purpose AI technology.

gr_app.py

@@ -0,0 +1,143 @@
+import gradio as gr
+import re
+from difflib import Differ
+from src.translation_agent.utils import *
+
+LANGUAGES = {
+    'English': 'English',
+    'Español': 'Spanish',
+    'Français': 'French',
+    'Deutsch': 'German',
+    'Italiano': 'Italian',
+    'Português': 'Portuguese',
+    'Русский': 'Russian',
+    '中文': 'Chinese',
+    '日本語': 'Japanese',
+    '한국어': 'Korean',
+    'العربية': 'Arabic',
+    'हिन्दी': 'Hindi',
+}
+
+
+def diff_texts(text1, text2, lang):
+    d = Differ()
+    ic(lang)
+    if lang == '中文':
+        return [
+            (token[2:], token[0] if token[0] != " " else None)
+            for token in d.compare(text1, text2)
+            if token[0] in ["+", " "]
+        ]
+    else:
+        words1 = re.findall(r'\S+|\s+', text1)
+        words2 = re.findall(r'\S+|\s+', text2)
+
+        return [
+            (token[2:], token[0] if token[0] != " " else None)
+            for token in d.compare(words1, words2)
+            if token[0] in ["+", " "]
+        ]
+    
+
+def translate_text(source_lang, target_lang, source_text, country, max_tokens=MAX_TOKENS_PER_CHUNK):
+    num_tokens_in_text = num_tokens_in_string(source_text)
+
+    ic(num_tokens_in_text)
+
+    if num_tokens_in_text < max_tokens:
+        ic("Translating text as single chunk")
+
+        #Note: use yield from B() if put yield in function B()
+        translation_1 = one_chunk_initial_translation(
+            source_lang, target_lang, source_text
+        )
+        yield translation_1, None, None
+
+        reflection = one_chunk_reflect_on_translation(
+            source_lang, target_lang, source_text, translation_1, country
+        )
+        yield translation_1, reflection, None
+
+        translation_2 = one_chunk_improve_translation(
+            source_lang, target_lang, source_text, translation_1, reflection
+        )
+        translation_diff = diff_texts(translation_1, translation_2, target_lang)
+        yield translation_1, reflection, translation_diff
+
+    else:
+        ic("Translating text as multiple chunks")
+
+        token_size = calculate_chunk_size(
+            token_count=num_tokens_in_text, token_limit=max_tokens
+        )
+
+        ic(token_size)
+
+        text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
+            model_name = "gpt-4",
+            chunk_size=token_size,
+            chunk_overlap=0,
+        )
+
+        source_text_chunks = text_splitter.split_text(source_text)
+        
+        translation_1_chunks = multichunk_initial_translation(
+            source_lang, target_lang, source_text_chunks
+        )
+        ic(translation_1_chunks)
+        translation_1 = "".join(translation_1_chunks)
+        yield translation_1, None, None
+
+        reflection_chunks = multichunk_reflect_on_translation(
+            source_lang,
+            target_lang,
+            source_text_chunks,
+            translation_1_chunks,
+            country,
+        )
+        ic(reflection_chunks)
+        reflection = "".join(reflection_chunks)
+        yield translation_1, reflection, None
+
+        translation_2_chunks = multichunk_improve_translation(
+            source_lang,
+            target_lang,
+            source_text_chunks,
+            translation_1_chunks,
+            reflection_chunks,
+        )
+        ic(translation_2_chunks)
+        translation_2 = "".join(translation_2_chunks)
+        translation_diff = diff_texts(translation_1, translation_2, target_lang)
+        
+        yield translation_1, reflection, translation_diff
+        
+
+
+def update_ui(translation_1, reflection, translation_diff):
+    return gr.update(value=translation_1), gr.update(value=reflection), gr.update(value=translation_diff)
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Andrew Ng's Translation Agent ")
+    with gr.Row():
+        source_lang = gr.Dropdown(choices=list(LANGUAGES.keys()), value='English', label="Source Language")
+        target_lang = gr.Dropdown(choices=list(LANGUAGES.keys()), value='中文', label="Target Language")
+        country = gr.Textbox(label="Country (for target language)")
+    source_text = gr.Textbox(label="Source Text", lines=5, show_copy_button=True)
+
+    btn = gr.Button("Translate")
+
+    with gr.Row():
+        translation_1 = gr.Textbox(label="Initial Translation", lines=3)
+        reflection = gr.Textbox(label="Reflection", lines=3)
+
+    translation_diff = gr.HighlightedText  (label="Final Translation",
+                                            combine_adjacent=True,
+                                            show_legend=True,
+                                            color_map={"+": "red"})
+    #translation = gr.Textbox(label="Final Translation", lines=5, show_copy_button=True)
+
+    btn.click(translate_text, inputs=[source_lang, target_lang, source_text, country], outputs=[translation_1, reflection, translation_diff], queue=True)
+    btn.click(update_ui, inputs=[translation_1, reflection, translation_diff], outputs=[translation_1, reflection, translation_diff], queue=True)
+
+demo.launch()

pyproject.toml

@@ -0,0 +1,121 @@
+[tool.poetry]
+name = "translation-agent"
+version = "0.1.0"
+description = "Agentic workflow for machine translation using LLMs"
+authors = ["Andrew Ng <[email protected]>"]
+license = "MIT"
+readme = "README.md"
+package-mode = true
+packages = [{ include = "translation_agent", from = "src" }]
+repository = "https://github.com/andrewyng/translation-agent"
+keywords = ["translation", "agents", "LLM", "machine translation"]
+
+
+[tool.poetry.dependencies]
+python = "^3.9"
+openai = "^1.28.1"
+tiktoken = "^0.6.0"
+joblib = "^1.4.2"
+pysrt = "^1.1.2"
+icecream = "^2.1.3"
+langchain-text-splitters = "^0.0.1"
+python-dotenv = "^1.0.1"
+gradio = "^4.36.1"
+
+[tool.poetry.group.dev]
+optional = true
+
+[tool.poetry.group.dev.dependencies]
+black = "^24.4.2"
+flake8 = "^7.0.0"
+pyright = "^1.1.362"
+pre-commit = "^3.7.1"
+ruff = "^0.4.4"
+
+[tool.poetry.group.test]
+optional = true
+
+[tool.poetry.group.test.dependencies]
+pytest = "^8.2.0"
+mypy = "^1.10.0"
+pytest-mock = "^3.14.0"
+
+[tool.poetry.group.eval]
+optional = true
+
+[tool.poetry.group.eval.dependencies]
+nltk = "^3.8.1"
+sacrebleu = "^2.4.2"
+google-cloud-translate = "^3.15.3"
+deepl = "^1.18.0"
+numpy = "^1.26.4"
+scipy = "^1.13.0"
+gradio = "^4.31.5"
+requests = "^2.32.3"
+beautifulsoup4 = "^4.12.3"
+sentencepiece = "^0.2.0"
+
+
+[[tool.poetry.source]]
+name = "pytorch"
+url = "https://download.pytorch.org/whl/nightly/rocm6.0"
+priority = "supplemental"
+
+[tool.ruff]
+# Set the maximum line length to 79.
+line-length = 79
+indent-width = 4
+exclude = [".venv", ".env", ".git", "tests", "eval"]
+
+[tool.ruff.lint]
+# Add the `line-too-long` rule to the enforced rule set. By default, Ruff omits rules that
+# overlap with the use of a formatter, like Black, but we can override this behavior by
+# explicitly adding the rule.
+extend-select = [
+    "B",   # flake8-bugbear
+    "C4",  # flake8-comprehensions
+    "ERA", # flake8-eradicate/eradicate
+    "I",   # isort
+    "N",   # pep8-naming
+    "PIE", # flake8-pie
+    "PGH", # pygrep
+    "RUF", # ruff checks
+    "SIM", # flake8-simplify
+    # "T20", # flake8-print
+    "TCH", # flake8-type-checking
+    "TID", # flake8-tidy-imports
+    "UP",  # pyupgrade
+]
+fixable = ["ALL"]
+ignore = ["SIM117"]
+
+[tool.ruff.lint.isort]
+force-single-line = true
+lines-after-imports = 2
+known-first-party = ["translation-agent"]
+
+[tool.ruff.lint.per-file-ignores]
+"**/__init__.py" = ["E402", "F401"]
+"**/{tests,docs,tools}/*" = ["E402"]
+
+
+[tool.mypy]
+files = "src, tests"
+mypy_path = "src"
+namespace_packages = true
+explicit_package_bases = true
+show_error_codes = true
+strict = true
+enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]
+exclude = ["tests"]
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+line-ending = "auto"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

requirements.txt

@@ -0,0 +1,92 @@
+aiofiles==23.2.1 ; python_version >= "3.9" and python_version < "4.0"
+altair==5.3.0 ; python_version >= "3.9" and python_version < "4.0"
+annotated-types==0.7.0 ; python_version >= "3.9" and python_version < "4.0"
+anyio==4.4.0 ; python_version >= "3.9" and python_version < "4.0"
+asttokens==2.4.1 ; python_version >= "3.9" and python_version < "4.0"
+attrs==23.2.0 ; python_version >= "3.9" and python_version < "4.0"
+certifi==2024.6.2 ; python_version >= "3.9" and python_version < "4.0"
+chardet==5.2.0 ; python_version >= "3.9" and python_version < "4.0"
+charset-normalizer==3.3.2 ; python_version >= "3.9" and python_version < "4.0"
+click==8.1.7 ; python_version >= "3.9" and python_version < "4.0"
+colorama==0.4.6 ; python_version >= "3.9" and python_version < "4.0"
+contourpy==1.2.1 ; python_version >= "3.9" and python_version < "4.0"
+cycler==0.12.1 ; python_version >= "3.9" and python_version < "4.0"
+distro==1.9.0 ; python_version >= "3.9" and python_version < "4.0"
+dnspython==2.6.1 ; python_version >= "3.9" and python_version < "4.0"
+email-validator==2.1.1 ; python_version >= "3.9" and python_version < "4.0"
+exceptiongroup==1.2.1 ; python_version >= "3.9" and python_version < "3.11"
+executing==2.0.1 ; python_version >= "3.9" and python_version < "4.0"
+fastapi-cli==0.0.4 ; python_version >= "3.9" and python_version < "4.0"
+fastapi==0.111.0 ; python_version >= "3.9" and python_version < "4.0"
+ffmpy==0.3.2 ; python_version >= "3.9" and python_version < "4.0"
+filelock==3.15.1 ; python_version >= "3.9" and python_version < "4.0"
+fonttools==4.53.0 ; python_version >= "3.9" and python_version < "4.0"
+fsspec==2024.6.0 ; python_version >= "3.9" and python_version < "4.0"
+gradio-client==1.0.1 ; python_version >= "3.9" and python_version < "4.0"
+gradio==4.36.1 ; python_version >= "3.9" and python_version < "4.0"
+h11==0.14.0 ; python_version >= "3.9" and python_version < "4.0"
+httpcore==1.0.5 ; python_version >= "3.9" and python_version < "4.0"
+httptools==0.6.1 ; python_version >= "3.9" and python_version < "4.0"
+httpx==0.27.0 ; python_version >= "3.9" and python_version < "4.0"
+huggingface-hub==0.23.4 ; python_version >= "3.9" and python_version < "4.0"
+icecream==2.1.3 ; python_version >= "3.9" and python_version < "4.0"
+idna==3.7 ; python_version >= "3.9" and python_version < "4.0"
+importlib-resources==6.4.0 ; python_version >= "3.9" and python_version < "4.0"
+jinja2==3.1.4 ; python_version >= "3.9" and python_version < "4.0"
+joblib==1.4.2 ; python_version >= "3.9" and python_version < "4.0"
+jsonpatch==1.33 ; python_version >= "3.9" and python_version < "4.0"
+jsonpointer==3.0.0 ; python_version >= "3.9" and python_version < "4.0"
+jsonschema-specifications==2023.12.1 ; python_version >= "3.9" and python_version < "4.0"
+jsonschema==4.22.0 ; python_version >= "3.9" and python_version < "4.0"
+kiwisolver==1.4.5 ; python_version >= "3.9" and python_version < "4.0"
+langchain-core==0.1.52 ; python_version >= "3.9" and python_version < "4.0"
+langchain-text-splitters==0.0.1 ; python_version >= "3.9" and python_version < "4.0"
+langsmith==0.1.77 ; python_version >= "3.9" and python_version < "4.0"
+markdown-it-py==3.0.0 ; python_version >= "3.9" and python_version < "4.0"
+markupsafe==2.1.5 ; python_version >= "3.9" and python_version < "4.0"
+matplotlib==3.9.0 ; python_version >= "3.9" and python_version < "4.0"
+mdurl==0.1.2 ; python_version >= "3.9" and python_version < "4.0"
+numpy==1.26.4 ; python_version >= "3.9" and python_version < "4.0"
+openai==1.34.0 ; python_version >= "3.9" and python_version < "4.0"
+orjson==3.10.5 ; python_version >= "3.9" and python_version < "4.0"
+packaging==23.2 ; python_version >= "3.9" and python_version < "4.0"
+pandas==2.2.2 ; python_version >= "3.9" and python_version < "4.0"
+pillow==10.3.0 ; python_version >= "3.9" and python_version < "4.0"
+pydantic-core==2.18.4 ; python_version >= "3.9" and python_version < "4.0"
+pydantic==2.7.4 ; python_version >= "3.9" and python_version < "4.0"
+pydub==0.25.1 ; python_version >= "3.9" and python_version < "4.0"
+pygments==2.18.0 ; python_version >= "3.9" and python_version < "4.0"
+pyparsing==3.1.2 ; python_version >= "3.9" and python_version < "4.0"
+pysrt==1.1.2 ; python_version >= "3.9" and python_version < "4.0"
+python-dateutil==2.9.0.post0 ; python_version >= "3.9" and python_version < "4.0"
+python-dotenv==1.0.1 ; python_version >= "3.9" and python_version < "4.0"
+python-multipart==0.0.9 ; python_version >= "3.9" and python_version < "4.0"
+pytz==2024.1 ; python_version >= "3.9" and python_version < "4.0"
+pyyaml==6.0.1 ; python_version >= "3.9" and python_version < "4.0"
+referencing==0.35.1 ; python_version >= "3.9" and python_version < "4.0"
+regex==2024.5.15 ; python_version >= "3.9" and python_version < "4.0"
+requests==2.32.3 ; python_version >= "3.9" and python_version < "4.0"
+rich==13.7.1 ; python_version >= "3.9" and python_version < "4.0"
+rpds-py==0.18.1 ; python_version >= "3.9" and python_version < "4.0"
+ruff==0.4.9 ; python_version >= "3.9" and python_version < "4.0" and sys_platform != "emscripten"
+semantic-version==2.10.0 ; python_version >= "3.9" and python_version < "4.0"
+shellingham==1.5.4 ; python_version >= "3.9" and python_version < "4.0"
+six==1.16.0 ; python_version >= "3.9" and python_version < "4.0"
+sniffio==1.3.1 ; python_version >= "3.9" and python_version < "4.0"
+starlette==0.37.2 ; python_version >= "3.9" and python_version < "4.0"
+tenacity==8.3.0 ; python_version >= "3.9" and python_version < "4.0"
+tiktoken==0.6.0 ; python_version >= "3.9" and python_version < "4.0"
+tomlkit==0.12.0 ; python_version >= "3.9" and python_version < "4.0"
+toolz==0.12.1 ; python_version >= "3.9" and python_version < "4.0"
+tqdm==4.66.4 ; python_version >= "3.9" and python_version < "4.0"
+typer==0.12.3 ; python_version >= "3.9" and python_version < "4.0"
+typing-extensions==4.12.2 ; python_version >= "3.9" and python_version < "4.0"
+tzdata==2024.1 ; python_version >= "3.9" and python_version < "4.0"
+ujson==5.10.0 ; python_version >= "3.9" and python_version < "4.0"
+urllib3==2.2.1 ; python_version >= "3.9" and python_version < "4.0"
+uvicorn==0.30.1 ; python_version >= "3.9" and python_version < "4.0" and sys_platform != "emscripten"
+uvicorn[standard]==0.30.1 ; python_version >= "3.9" and python_version < "4.0"
+uvloop==0.19.0 ; (sys_platform != "win32" and sys_platform != "cygwin") and platform_python_implementation != "PyPy" and python_version >= "3.9" and python_version < "4.0"
+watchfiles==0.22.0 ; python_version >= "3.9" and python_version < "4.0"
+websockets==11.0.3 ; python_version >= "3.9" and python_version < "4.0"
+zipp==3.19.2 ; python_version >= "3.9" and python_version < "3.10"

src/translation_agent/__init__.py

@@ -0,0 +1 @@
+from .utils import translate

src/translation_agent/utils.py

@@ -0,0 +1,689 @@
+import os
+from typing import List
+from typing import Union
+
+import openai
+import tiktoken
+from dotenv import load_dotenv
+from icecream import ic
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+
+load_dotenv()  # read local .env file
+model = os.getenv("OPENAI_MODEL") or "gpt-4-turbo"
+client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"), base_url=os.getenv("OPENAI_BASE_URL"))
+
+
+MAX_TOKENS_PER_CHUNK = (
+    1000  # if text is more than this many tokens, we'll break it up into
+)
+# discrete chunks to translate one chunk at a time
+
+
+def get_completion(
+    prompt: str,
+    system_message: str = "You are a helpful assistant.",
+    model: str = model,
+    temperature: float = 0.3,
+    json_mode: bool = False,
+) -> Union[str, dict]:
+    """
+        Generate a completion using the OpenAI API.
+
+    Args:
+        prompt (str): The user's prompt or query.
+        system_message (str, optional): The system message to set the context for the assistant.
+            Defaults to "You are a helpful assistant.".
+        model (str, optional): The name of the OpenAI model to use for generating the completion.
+            Defaults to "gpt-4-turbo".
+        temperature (float, optional): The sampling temperature for controlling the randomness of the generated text.
+            Defaults to 0.3.
+        json_mode (bool, optional): Whether to return the response in JSON format.
+            Defaults to False.
+
+    Returns:
+        Union[str, dict]: The generated completion.
+            If json_mode is True, returns the complete API response as a dictionary.
+            If json_mode is False, returns the generated text as a string.
+    """
+
+    if json_mode:
+        response = client.chat.completions.create(
+            model=model,
+            temperature=temperature,
+            top_p=1,
+            response_format={"type": "json_object"},
+            messages=[
+                {"role": "system", "content": system_message},
+                {"role": "user", "content": prompt},
+            ],
+        )
+        return response.choices[0].message.content
+    else:
+        response = client.chat.completions.create(
+            model=model,
+            temperature=temperature,
+            top_p=1,
+            messages=[
+                {"role": "system", "content": system_message},
+                {"role": "user", "content": prompt},
+            ],
+        )
+        return response.choices[0].message.content
+
+
+def one_chunk_initial_translation(
+    source_lang: str, target_lang: str, source_text: str
+) -> str:
+    """
+    Translate the entire text as one chunk using an LLM.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language for translation.
+        source_text (str): The text to be translated.
+
+    Returns:
+        str: The translated text.
+    """
+
+    system_message = f"You are an expert linguist, specializing in translation 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}
+
+{target_lang}:"""
+
+    prompt = translation_prompt.format(source_text=source_text)
+
+    translation = get_completion(prompt, system_message=system_message)
+
+    return translation
+
+
+def one_chunk_reflect_on_translation(
+    source_lang: str,
+    target_lang: str,
+    source_text: str,
+    translation_1: str,
+    country: str = "",
+) -> str:
+    """
+    Use an LLM to reflect on the translation, treating the entire text as one chunk.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language of the translation.
+        source_text (str): The original text in the source language.
+        translation_1 (str): The initial translation of the source text.
+        country (str): Country specified for target language.
+
+    Returns:
+        str: The LLM's reflection on the translation, providing constructive criticism and suggestions for improvement.
+    """
+
+    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."
+
+    if country != "":
+        reflection_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></SOURCE_TEXT> and <TRANSLATION></TRANSLATION>, are as follows:
+
+<SOURCE_TEXT>
+{source_text}
+</SOURCE_TEXT>
+
+<TRANSLATION>
+{translation_1}
+</TRANSLATION>
+
+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.
+Output only the suggestions and nothing else."""
+
+    else:
+        reflection_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 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>
+{translation_1}
+</TRANSLATION>
+
+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.
+Output only the suggestions and nothing else."""
+
+    prompt = reflection_prompt.format(
+        source_lang=source_lang,
+        target_lang=target_lang,
+        source_text=source_text,
+        translation_1=translation_1,
+    )
+    reflection = get_completion(prompt, system_message=system_message)
+    return reflection
+
+
+def one_chunk_improve_translation(
+    source_lang: str,
+    target_lang: str,
+    source_text: str,
+    translation_1: str,
+    reflection: str,
+) -> str:
+    """
+    Use the reflection to improve the translation, treating the entire text as one chunk.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language for the translation.
+        source_text (str): The original text in the source language.
+        translation_1 (str): The initial translation of the source text.
+        reflection (str): Expert suggestions and constructive criticism for improving the translation.
+
+    Returns:
+        str: The improved translation based on the expert suggestions.
+    """
+
+    system_message = f"You are an expert linguist, specializing in translation editing from {source_lang} to {target_lang}."
+
+    prompt = f"""Your task is to carefully read, then edit, a translation from {source_lang} to {target_lang}, taking into
+account a list of expert suggestions and constructive criticisms.
+
+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>
+{source_text}
+</SOURCE_TEXT>
+
+<TRANSLATION>
+{translation_1}
+</TRANSLATION>
+
+<EXPERT_SUGGESTIONS>
+{reflection}
+</EXPERT_SUGGESTIONS>
+
+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 {target_lang} 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 only the new translation and nothing else."""
+
+    translation_2 = get_completion(prompt, system_message)
+
+    return translation_2
+
+
+def one_chunk_translate_text(
+    source_lang: str, target_lang: str, source_text: str, country: str = ""
+) -> str:
+    """
+    Translate a single chunk of text from the source language to the target language.
+
+    This function performs a two-step translation process:
+    1. Get an initial translation of the source text.
+    2. Reflect on the initial translation and generate an improved translation.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language for the translation.
+        source_text (str): The text to be translated.
+        country (str): Country specified for target language.
+    Returns:
+        str: The improved translation of the source text.
+    """
+    translation_1 = one_chunk_initial_translation(
+        source_lang, target_lang, source_text
+    )
+
+    reflection = one_chunk_reflect_on_translation(
+        source_lang, target_lang, source_text, translation_1, country
+    )
+    translation_2 = one_chunk_improve_translation(
+        source_lang, target_lang, source_text, translation_1, reflection
+    )
+
+    return translation_2
+
+
+def num_tokens_in_string(
+    input_str: str, encoding_name: str = "cl100k_base"
+) -> int:
+    """
+    Calculate the number of tokens in a given string using a specified encoding.
+
+    Args:
+        str (str): The input string to be tokenized.
+        encoding_name (str, optional): The name of the encoding to use. Defaults to "cl100k_base",
+            which is the most commonly used encoder (used by GPT-4).
+
+    Returns:
+        int: The number of tokens in the input string.
+
+    Example:
+        >>> text = "Hello, how are you?"
+        >>> num_tokens = num_tokens_in_string(text)
+        >>> print(num_tokens)
+        5
+    """
+    encoding = tiktoken.get_encoding(encoding_name)
+    num_tokens = len(encoding.encode(input_str))
+    return num_tokens
+
+
+def multichunk_initial_translation(
+    source_lang: str, target_lang: str, source_text_chunks: List[str]
+) -> List[str]:
+    """
+    Translate a text in multiple chunks from the source language to the target language.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language for translation.
+        source_text_chunks (List[str]): A list of text chunks to be translated.
+
+    Returns:
+        List[str]: A list of translated text chunks.
+    """
+
+    system_message = f"You are an expert linguist, specializing in translation from {source_lang} to {target_lang}."
+
+    translation_prompt = """Your task is provide a professional translation from {source_lang} to {target_lang} of PART of a text.
+
+The source text is below, delimited by XML tags <SOURCE_TEXT> and </SOURCE_TEXT>. Translate only the part within the source text
+delimited by <TRANSLATE_THIS> and </TRANSLATE_THIS>. You can use the rest of the source text as context, but do not translate any
+of the other text. Do not output anything other than the translation of the indicated part of the text.
+
+<SOURCE_TEXT>
+{tagged_text}
+</SOURCE_TEXT>
+
+To reiterate, you should translate only this part of the text, shown here again between <TRANSLATE_THIS> and </TRANSLATE_THIS>:
+<TRANSLATE_THIS>
+{chunk_to_translate}
+</TRANSLATE_THIS>
+
+Output only the translation of the portion you are asked to translate, and nothing else.
+"""
+
+    translation_chunks = []
+    for i in range(len(source_text_chunks)):
+        # Will translate chunk i
+        tagged_text = (
+            "".join(source_text_chunks[0:i])
+            + "<TRANSLATE_THIS>"
+            + source_text_chunks[i]
+            + "</TRANSLATE_THIS>"
+            + "".join(source_text_chunks[i + 1 :])
+        )
+
+        prompt = translation_prompt.format(
+            source_lang=source_lang,
+            target_lang=target_lang,
+            tagged_text=tagged_text,
+            chunk_to_translate=source_text_chunks[i],
+        )
+
+        translation = get_completion(prompt, system_message=system_message)
+        translation_chunks.append(translation)
+
+    return translation_chunks
+
+
+def multichunk_reflect_on_translation(
+    source_lang: str,
+    target_lang: str,
+    source_text_chunks: List[str],
+    translation_1_chunks: List[str],
+    country: str = "",
+) -> List[str]:
+    """
+    Provides constructive criticism and suggestions for improving a partial translation.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language of the translation.
+        source_text_chunks (List[str]): The source text divided into chunks.
+        translation_1_chunks (List[str]): The translated chunks corresponding to the source text chunks.
+        country (str): Country specified for target language.
+
+    Returns:
+        List[str]: A list of reflections containing suggestions for improving each translated chunk.
+    """
+
+    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."
+
+    if country != "":
+        reflection_prompt = """Your task is to carefully read a source text and part of a translation of that text from {source_lang} to {target_lang}, and then give constructive criticism and helpful suggestions for improving the translation.
+The final style and tone of the translation should match the style of {target_lang} colloquially spoken in {country}.
+
+The source text is below, delimited by XML tags <SOURCE_TEXT> and </SOURCE_TEXT>, and the part that has been translated
+is delimited by <TRANSLATE_THIS> and </TRANSLATE_THIS> within the source text. You can use the rest of the source text
+as context for critiquing the translated part.
+
+<SOURCE_TEXT>
+{tagged_text}
+</SOURCE_TEXT>
+
+To reiterate, only part of the text is being translated, shown here again between <TRANSLATE_THIS> and </TRANSLATE_THIS>:
+<TRANSLATE_THIS>
+{chunk_to_translate}
+</TRANSLATE_THIS>
+
+The translation of the indicated part, delimited below by <TRANSLATION> and </TRANSLATION>, is as follows:
+<TRANSLATION>
+{translation_1_chunk}
+</TRANSLATION>
+
+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.
+Output only the suggestions and nothing else."""
+
+    else:
+        reflection_prompt = """Your task is to carefully read a source text and part of a translation of that text from {source_lang} to {target_lang}, and then give constructive criticism and helpful suggestions for improving the translation.
+
+The source text is below, delimited by XML tags <SOURCE_TEXT> and </SOURCE_TEXT>, and the part that has been translated
+is delimited by <TRANSLATE_THIS> and </TRANSLATE_THIS> within the source text. You can use the rest of the source text
+as context for critiquing the translated part.
+
+<SOURCE_TEXT>
+{tagged_text}
+</SOURCE_TEXT>
+
+To reiterate, only part of the text is being translated, shown here again between <TRANSLATE_THIS> and </TRANSLATE_THIS>:
+<TRANSLATE_THIS>
+{chunk_to_translate}
+</TRANSLATE_THIS>
+
+The translation of the indicated part, delimited below by <TRANSLATION> and </TRANSLATION>, is as follows:
+<TRANSLATION>
+{translation_1_chunk}
+</TRANSLATION>
+
+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.
+Output only the suggestions and nothing else."""
+
+    reflection_chunks = []
+    for i in range(len(source_text_chunks)):
+        # Will translate chunk i
+        tagged_text = (
+            "".join(source_text_chunks[0:i])
+            + "<TRANSLATE_THIS>"
+            + source_text_chunks[i]
+            + "</TRANSLATE_THIS>"
+            + "".join(source_text_chunks[i + 1 :])
+        )
+        if country != "":
+            prompt = reflection_prompt.format(
+                source_lang=source_lang,
+                target_lang=target_lang,
+                tagged_text=tagged_text,
+                chunk_to_translate=source_text_chunks[i],
+                translation_1_chunk=translation_1_chunks[i],
+                country=country,
+            )
+        else:
+            prompt = reflection_prompt.format(
+                source_lang=source_lang,
+                target_lang=target_lang,
+                tagged_text=tagged_text,
+                chunk_to_translate=source_text_chunks[i],
+                translation_1_chunk=translation_1_chunks[i],
+            )
+
+        reflection = get_completion(prompt, system_message=system_message)
+        reflection_chunks.append(reflection)
+
+    return reflection_chunks
+
+
+def multichunk_improve_translation(
+    source_lang: str,
+    target_lang: str,
+    source_text_chunks: List[str],
+    translation_1_chunks: List[str],
+    reflection_chunks: List[str],
+) -> List[str]:
+    """
+    Improves the translation of a text from source language to target language by considering expert suggestions.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language for translation.
+        source_text_chunks (List[str]): The source text divided into chunks.
+        translation_1_chunks (List[str]): The initial translation of each chunk.
+        reflection_chunks (List[str]): Expert suggestions for improving each translated chunk.
+
+    Returns:
+        List[str]: The improved translation of each chunk.
+    """
+
+    system_message = f"You are an expert linguist, specializing in translation editing from {source_lang} to {target_lang}."
+
+    improvement_prompt = """Your task is to carefully read, then improve, a translation from {source_lang} to {target_lang}, taking into
+account a set of expert suggestions and constructive critisms. Below, the source text, initial translation, and expert suggestions are provided.
+
+The source text is below, delimited by XML tags <SOURCE_TEXT> and </SOURCE_TEXT>, and the part that has been translated
+is delimited by <TRANSLATE_THIS> and </TRANSLATE_THIS> within the source text. You can use the rest of the source text
+as context, but need to provide a translation only of the part indicated by <TRANSLATE_THIS> and </TRANSLATE_THIS>.
+
+<SOURCE_TEXT>
+{tagged_text}
+</SOURCE_TEXT>
+
+To reiterate, only part of the text is being translated, shown here again between <TRANSLATE_THIS> and </TRANSLATE_THIS>:
+<TRANSLATE_THIS>
+{chunk_to_translate}
+</TRANSLATE_THIS>
+
+The translation of the indicated part, delimited below by <TRANSLATION> and </TRANSLATION>, is as follows:
+<TRANSLATION>
+{translation_1_chunk}
+</TRANSLATION>
+
+The expert translations of the indicated part, delimited below by <EXPERT_SUGGESTIONS> and </EXPERT_SUGGESTIONS>, is as follows:
+<EXPERT_SUGGESTIONS>
+{reflection_chunk}
+</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, or untranslated text),
+(ii) fluency (by applying {target_lang} 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 only the new translation of the indicated part and nothing else."""
+
+    translation_2_chunks = []
+    for i in range(len(source_text_chunks)):
+        # Will translate chunk i
+        tagged_text = (
+            "".join(source_text_chunks[0:i])
+            + "<TRANSLATE_THIS>"
+            + source_text_chunks[i]
+            + "</TRANSLATE_THIS>"
+            + "".join(source_text_chunks[i + 1 :])
+        )
+
+        prompt = improvement_prompt.format(
+            source_lang=source_lang,
+            target_lang=target_lang,
+            tagged_text=tagged_text,
+            chunk_to_translate=source_text_chunks[i],
+            translation_1_chunk=translation_1_chunks[i],
+            reflection_chunk=reflection_chunks[i],
+        )
+
+        translation_2 = get_completion(prompt, system_message=system_message)
+        translation_2_chunks.append(translation_2)
+
+    return translation_2_chunks
+
+
+def multichunk_translation(
+    source_lang, target_lang, source_text_chunks, country: str = ""
+):
+    """
+    Improves the translation of multiple text chunks based on the initial translation and reflection.
+
+    Args:
+        source_lang (str): The source language of the text chunks.
+        target_lang (str): The target language for translation.
+        source_text_chunks (List[str]): The list of source text chunks to be translated.
+        translation_1_chunks (List[str]): The list of initial translations for each source text chunk.
+        reflection_chunks (List[str]): The list of reflections on the initial translations.
+        country (str): Country specified for target language
+    Returns:
+        List[str]: The list of improved translations for each source text chunk.
+    """
+
+    translation_1_chunks = multichunk_initial_translation(
+        source_lang, target_lang, source_text_chunks
+    )
+
+    reflection_chunks = multichunk_reflect_on_translation(
+        source_lang,
+        target_lang,
+        source_text_chunks,
+        translation_1_chunks,
+        country,
+    )
+
+    translation_2_chunks = multichunk_improve_translation(
+        source_lang,
+        target_lang,
+        source_text_chunks,
+        translation_1_chunks,
+        reflection_chunks,
+    )
+
+    return translation_2_chunks
+
+
+def calculate_chunk_size(token_count: int, token_limit: int) -> int:
+    """
+    Calculate the chunk size based on the token count and token limit.
+
+    Args:
+        token_count (int): The total number of tokens.
+        token_limit (int): The maximum number of tokens allowed per chunk.
+
+    Returns:
+        int: The calculated chunk size.
+
+    Description:
+        This function calculates the chunk size based on the given token count and token limit.
+        If the token count is less than or equal to the token limit, the function returns the token count as the chunk size.
+        Otherwise, it calculates the number of chunks needed to accommodate all the tokens within the token limit.
+        The chunk size is determined by dividing the token limit by the number of chunks.
+        If there are remaining tokens after dividing the token count by the token limit,
+        the chunk size is adjusted by adding the remaining tokens divided by the number of chunks.
+
+    Example:
+        >>> calculate_chunk_size(1000, 500)
+        500
+        >>> calculate_chunk_size(1530, 500)
+        389
+        >>> calculate_chunk_size(2242, 500)
+        496
+    """
+
+    if token_count <= token_limit:
+        return token_count
+
+    num_chunks = (token_count + token_limit - 1) // token_limit
+    chunk_size = token_count // num_chunks
+
+    remaining_tokens = token_count % token_limit
+    if remaining_tokens > 0:
+        chunk_size += remaining_tokens // num_chunks
+
+    return chunk_size
+
+
+def translate(
+    source_lang,
+    target_lang,
+    source_text,
+    country,
+    max_tokens=MAX_TOKENS_PER_CHUNK,
+):
+    """Translate the source_text from source_lang to target_lang."""
+
+    num_tokens_in_text = num_tokens_in_string(source_text)
+
+    ic(num_tokens_in_text)
+
+    if num_tokens_in_text < max_tokens:
+        ic("Translating text as single chunk")
+
+        final_translation = one_chunk_translate_text(
+            source_lang, target_lang, source_text, country
+        )
+
+        return final_translation
+
+    else:
+        ic("Translating text as multiple chunks")
+
+        token_size = calculate_chunk_size(
+            token_count=num_tokens_in_text, token_limit=max_tokens
+        )
+
+        ic(token_size)
+
+        text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
+            model_name="gpt-4",
+            chunk_size=token_size,
+            chunk_overlap=0,
+        )
+
+        source_text_chunks = text_splitter.split_text(source_text)
+
+        translation_2_chunks = multichunk_translation(
+            source_lang, target_lang, source_text_chunks, country
+        )
+
+        return "".join(translation_2_chunks)

tests/test_agent.py

@@ -0,0 +1,289 @@
+import json
+import os
+from unittest.mock import patch
+
+import openai
+import pytest
+from dotenv import load_dotenv
+
+# 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()
+
+client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
+
+def test_get_completion_json_mode_api_call():
+    # Set up the test data
+    prompt = "What is the capital of France in json?"
+    system_message = "You are a helpful assistant."
+    model = "gpt-4-turbo"
+    temperature = 0.3
+    json_mode = True
+
+    # Call the function with JSON_mode=True
+    result = get_completion(
+        prompt, system_message, model, temperature, json_mode
+    )
+
+    # Assert that the result is not None
+    assert result is not None
+
+    # Assert that it can be transformed to dictionary (json)
+    assert isinstance(json.loads(result), dict)
+
+
+def test_get_completion_non_json_mode_api_call():
+    # Set up the test data
+    prompt = "What is the capital of France?"
+    system_message = "You are a helpful assistant."
+    model = "gpt-4-turbo"
+    temperature = 0.3
+    json_mode = False
+
+    # Call the function with JSON_mode=False
+    result = get_completion(
+        prompt, system_message, model, temperature, json_mode
+    )
+
+    # Assert that the result is not None
+    assert result is not None
+
+    # Assert that the result has the expected response format
+    assert isinstance(result, str)
+
+
+def test_one_chunk_initial_translation():
+    # Define test data
+    source_lang = "English"
+    target_lang = "Spanish"
+    source_text = "Hello, how are you?"
+    expected_translation = "Hola, ¿cómo estás?"
+
+    # Mock the get_completion_content function
+    with patch(
+        "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(
+            source_lang, target_lang, source_text
+        )
+
+        # Assert the expected translation is returned
+        assert translation == expected_translation
+
+        # Assert the get_completion_content function was called with the correct arguments
+        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}
+
+{target_lang}:"""
+
+        mock_get_completion.assert_called_once_with(
+            expected_prompt, system_message=expected_system_message
+        )
+
+
+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."
+    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.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, 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 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></SOURCE_TEXT> and <TRANSLATION></TRANSLATION>, are as follows:
+
+<SOURCE_TEXT>
+{source_text}
+</SOURCE_TEXT>
+
+<TRANSLATION>
+{translation_1}
+</TRANSLATION>
+
+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.
+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
+        )
+
+
+@pytest.fixture
+def example_data():
+    return {
+        "source_lang": "English",
+        "target_lang": "Spanish",
+        "source_text": "This is a sample source text.",
+        "translation_1": "Esta es una traducción de ejemplo.",
+        "reflection": "The translation is accurate but could be more fluent.",
+    }
+
+
+@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.return_value = (
+        "Esta es una traducción de ejemplo mejorada."
+    )
+
+    # Call the function with the example data
+    result = one_chunk_improve_translation(
+        example_data["source_lang"],
+        example_data["target_lang"],
+        example_data["source_text"],
+        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 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, 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>
+{example_data["source_text"]}
+</SOURCE_TEXT>
+
+<TRANSLATION>
+{example_data["translation_1"]}
+</TRANSLATION>
+
+<EXPERT_SUGGESTIONS>
+{example_data["reflection"]}
+</EXPERT_SUGGESTIONS>
+
+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 only the new translation and nothing else."""
+
+    expected_system_message = f"You are an expert linguist, specializing in translation editing from English to Spanish."
+
+    mock_get_completion.assert_called_once_with(
+        expected_prompt, expected_system_message
+    )
+
+
+def test_one_chunk_translate_text(mocker):
+    # Define test data
+    source_lang = "English"
+    target_lang = "Spanish"
+    country = "Mexico"
+    source_text = "Hello, how are you?"
+    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.utils.one_chunk_initial_translation",
+        return_value=translation_1,
+    )
+    mock_reflect_on_translation = mocker.patch(
+        "translation_agent.utils.one_chunk_reflect_on_translation",
+        return_value=reflection,
+    )
+    mock_improve_translation = mocker.patch(
+        "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, country
+    )
+
+    # Assert the expected result
+    assert result == translation2
+
+    # Assert that the helper functions were called with the correct arguments
+    mock_initial_translation.assert_called_once_with(
+        source_lang, target_lang, source_text
+    )
+    mock_reflect_on_translation.assert_called_once_with(
+        source_lang, target_lang, source_text, translation_1, country
+    )
+    mock_improve_translation.assert_called_once_with(
+        source_lang, target_lang, source_text, translation_1, reflection
+    )
+
+
+def test_num_tokens_in_string():
+    # Test case 1: Empty string
+    assert num_tokens_in_string("") == 0
+
+    # Test case 2: Simple string
+    assert num_tokens_in_string("Hello, world!") == 4
+
+    # Test case 3: String with special characters
+    assert (
+        num_tokens_in_string(
+            "This is a test string with special characters: !@#$%^&*()"
+        )
+        == 16
+    )
+
+    # Test case 4: String with non-ASCII characters
+    assert num_tokens_in_string("Héllò, wörld! 你好，世界！") == 17
+
+    # Test case 5: Long string
+    long_string = (
+        "Lorem ipsum dolor sit amet, consectetur adipiscing elit. " * 10
+    )
+    assert num_tokens_in_string(long_string) == 101
+
+    # Test case 6: Different encoding
+    assert (
+        num_tokens_in_string("Hello, world!", encoding_name="p50k_base") == 4
+    )