Added unit test and removed parameters

src/translator.py

@@ -3,8 +3,7 @@
 # from google.cloud import aiplatform
 import google.generativeai as genai
 
-# os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = "/Users/larissatyagi/Desktop/translator-service/translator-service-418821-9352d29e6139.json"
-# aiplatform.init(project='translator-service-418821', location='us-central1')
+
 os.environ['GOOGLE_API_KEY'] = 'AIzaSyA5nw5uJld70nkV-0D2C1gmhqo5ql9OdRw'  # Replace with your actual API key
 genai.configure(api_key=os.environ['GOOGLE_API_KEY'])
 
@@ -17,10 +16,6 @@ def get_translation(post: str) -> str:
                 +"please translate this to English. "
                 + f"Prompt: {post}")
     model = genai.GenerativeModel(model_name="gemini-pro")
-    parameters = {
-        "temperature": 0.7,  # Temperature controls the degree of randomness in token selection.
-        "max_output_tokens": 256,  # Token limit determines the maximum amount of text output.
-    }
 
     response = model.generate_content(context)
     return response.text
@@ -30,10 +25,6 @@ def get_translation(post: str) -> str:
 
 def get_language(post: str) -> str:
     model = genai.GenerativeModel(model_name="gemini-pro")
-    parameters = {
-        "temperature": 0.7,  # Temperature controls the degree of randomness in token selection.
-        "max_output_tokens": 256,  # Token limit determines the maximum amount of text output.
-    }
 
     context = f"Context = 'I want to know what language the prompt is', prompt = '{post}'"
     response = model.generate_content(context+post)

test/unit/test_translator.py

@@ -1,4 +1,4 @@
-# from src.translator import translate_content #,query_llm_robust
+from src.translator import translate_content #,query_llm_robust
 # from mock import patch
 # from vertexai.language_models import ChatModel, InputOutputTextPair
 # from ../src/translator import query_llm_robust
@@ -16,15 +16,20 @@
 #   assert response == "I don't understand your request"
 #   assert response is not None and response != ""
 
+### UNIT TESTING ###
+def test_chinese():
+    is_english, translated_content = translate_content("这是一条中文消息")
+    assert is_english == False
+    assert ("chinese" in translated_content or "Chinese" in translate_content)
 
-# def test_chinese():
-#     is_english, translated_content = translate_content("这是一条中文消息")
-#     assert is_english == False
-#     assert translated_content == "This is a Chinese message"
 
+def test_llm_normal_response():
+    is_english, translated_content = translate_content("to be or not to be")
+    assert is_english == True
+    assert ("to be or not to be"== translate_content)
+
 
-# def test_llm_normal_response():
-#     pass
-
-# def test_llm_gibberish_response():
-#     pass
+def test_llm_gibberish_response():
+    is_english, translated_content = translate_content("jfdjshghui bjfhsdufh")
+    assert ((is_english == False) or translated_content == "jfdjshghui bjfhsdufh")
+