adding supporting Amazon Bedrock

.env

@@ -0,0 +1,4 @@
+ACCESS_KEY="your aws ak"
+SECRET_KEY="your aws sk"
+OPENAI_API_KEY="your open api key"
+BEDROCK="True"

README.md

@@ -1,5 +1,12 @@
 # Translation Agent: Agentic translation using reflection workflow
 
+# adding supporting Amazon Bedrock   
+main changes:  
+src/translation_agent/utils.py  
++.env  
++test.py  
+
+
 This is a Python demonstration of a reflection agentic workflow for machine translation. The main steps are:
 1. Prompt an LLM to translate a text from `source_language` to `target_language`; 
 2. Have the LLM reflect on the translation to come up with constructive suggestions for improving it; 
@@ -32,20 +39,57 @@ To get started with `translation-agent`, follow these steps:
 pip install poetry 
 ```
 
-- A .env file with a OPENAI_API_KEY is required to run the workflow. See the .env.sample file as an example.
-```bash
-git clone https://github.com/andrewyng/translation-agent.git
-cd translation-agent
-poetry install
-poetry shell # activates virtual environment
+```bash  
+git clone https://github.com/shenshaoyong/translation-agent.git  
+cd translation-agent  
+poetry install  
+poetry shell   
+pip install -qU boto3 botocore  
+```
+
+- A .env file with a OPENAI_API_KEY / Amazon AKSK is required to run the workflow. See the .env.sample file as an example.  
+vi .env    
+```bash  
+ACCESS_KEY="your aws ak"
+SECRET_KEY="your aws sk"
+OPENAI_API_KEY="your open api key"
+BEDROCK="True"
+
 ```
-### Usage:
 
+### Usage:  
+vi test.py   
 ```python
-import translation_agent as ta
-source_lang, target_lang, country = "English", "Spanish", "Mexico"
-translation = ta.translate(source_lang, target_lang, source_text, country)
+import translation_agent as ta  
+source_text = "Today is a good day. Sunny, shine, breeze, blue sky"  
+source_lang, target_lang, country = "English", "Simplified Chinese", "China"  
+translation = ta.translate(source_lang, target_lang, source_text, country)  
+```
+
+run  
+```bash 
+python test.py 
 ```
+
+### result:  
+```text
+ic| num_tokens_in_text: 14  
+ic| 'Translating text as single chunk'  
+今天是个好天。阳光明媚，微风拂面，蓝天如洗。  
+1. 将"今天是个好天"改为"今天是个好日子"或"今天天气真好"，更符合口语表达习惯。  
+
+2. 考虑将"阳光明媚"改为"阳光灿烂"或"阳光普照"，以更贴近原文的"shine"。  
+
+3. "微风拂面"虽然是优美的表达，但可以简化为"微风轻拂"或"和风徐徐"，更接近原文的简洁风格。  
+
+4. "蓝天如洗"是很好的表达，但可以考虑改为"天空湛蓝"或"蓝天万里"，以更直接地对应原文的"blue sky"。  
+
+5. 可以考虑在句子之间添加逗号，使整体结构更加紧凑，例如："今天天气真好，阳光灿烂，和风徐徐，蓝天万里。"  
+
+6. 原文中的"Sunny, shine"有重复之意，翻译时可以合并处理，只保留一个相关描述即可。  
+今天天气真好，阳光灿烂，和风徐徐，蓝天万里。  
+```
+
 See examples/example_script.py for an example script to try out.
 
 ## License

src/translation_agent/utils.py

@@ -8,6 +8,9 @@
 from icecream import ic
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 
+import boto3
+import botocore
+import logging
 
 load_dotenv()  # read local .env file
 client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
@@ -44,6 +47,8 @@ def get_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 os.getenv("BEDROCK")=="True":
+        return get_completion_bedrock(prompt=prompt,system_message=system_message,model="anthropic.claude-3-5-sonnet-20240620-v1:0",temperature=0.3,json_mode=False)
 
     if json_mode:
         response = client.chat.completions.create(
@@ -69,6 +74,77 @@ def get_completion(
         )
         return response.choices[0].message.content
 
+def get_completion_bedrock(
+    prompt: str,
+    system_message: str = "You are a helpful assistant.",
+    model: str = "anthropic.claude-3-5-sonnet-20240620-v1:0",
+    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.
+    """
+    messages = [
+        {
+            "role": "user",
+            "content": [{
+                "text": prompt
+            }]
+        }
+    ]
+
+    bedrock_client = boto3.client(service_name='bedrock-runtime', 
+        region_name='us-east-1',
+        aws_access_key_id=os.getenv("ACCESS_KEY"),
+        aws_secret_access_key=os.getenv("SECRET_KEY"))
+
+    system_prompts = [{"text" : system_message}]
+
+    # Inference parameters to use.
+    top_k = 200
+
+    #Base inference parameters to use.
+    inference_config = {"temperature": temperature}
+    # Additional inference parameters to use.
+    additional_model_fields = {"top_k": top_k}
+    try:
+        # Send the message.
+        response = bedrock_client.converse(
+            modelId=model,
+            messages=messages,
+            system=system_prompts,
+            inferenceConfig=inference_config,
+            additionalModelRequestFields=additional_model_fields,
+        )
+
+        #print(response['output'])
+        #print(response['usage'])
+    except ClientError as err:
+        message = err.response['Error']['Message']
+        logger.error("A client error occurred: %s", message)
+        #print(f"A client error occured: {message}")
+
+    else:
+        pass
+        #print(f"Finished generating text with model {model}.")
+    print(response['output']['message']['content'][0]['text'])
+    return response['output']['message']['content'][0]['text']
 
 def one_chunk_initial_translation(
     source_lang: str, target_lang: str, source_text: str

test.py

@@ -0,0 +1,4 @@
+import translation_agent as ta
+source_text = "Today is a good day. Sunny, shine, breeze, blue sky"
+source_lang, target_lang, country,max_tokens = "English", "Simplified Chinese", "China", 100
+translation = ta.translate(source_lang, target_lang, source_text, country,max_tokens=max_tokens)