updated gemini tests

Author: firojalam

llmebench/models/Gemini.py

@@ -6,9 +6,8 @@
 
 import vertexai
 import vertexai.preview.generative_models as generative_models
-from google.oauth2 import service_account
 from vertexai.generative_models import FinishReason, GenerativeModel, Part
-
+from google.oauth2 import service_account
 from llmebench.models.model_base import ModelBase
 
 
@@ -56,11 +55,11 @@ def __init__(
         project_id=None,
         model_name=None,
         location=None,
-        credentials_path=None,  # path to JSON file
-        credentials_info=None,  # dict or JSON string
+        credentials_path=None,      # path to JSON file
+        credentials_info=None,      # dict or JSON string
         timeout=20,
         temperature=0,
-        tolerance=1e-7,
+        tolerance = 1e-7,
         top_p=0.95,
         max_tokens=2000,
         **kwargs,
@@ -74,34 +73,28 @@ def __init__(
         if credentials_info:
             if isinstance(credentials_info, str):
                 credentials_info = json.loads(credentials_info)
-            self.credentials = service_account.Credentials.from_service_account_info(
-                credentials_info
-            )
+            self.credentials = service_account.Credentials.from_service_account_info(credentials_info)
         # 2. Else, load from path (arg or env)
         elif credentials_path or os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
             path = credentials_path or os.getenv("GOOGLE_APPLICATION_CREDENTIALS")
             with open(path, "r") as f:
                 info = json.load(f)
-            self.credentials = service_account.Credentials.from_service_account_info(
-                info
-            )
+            self.credentials = service_account.Credentials.from_service_account_info(info)
+        elif os.getenv("GOOGLE_APPLICATION_CREDENTIALS") is not None:
+            os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.getenv("GOOGLE_APPLICATION_CREDENTIALS")
         # 3. Else, None: will fall back to ADC (Application Default Credentials)
 
         if not self.project_id:
-            raise Exception(
-                "PROJECT_ID must be set (argument or `GOOGLE_PROJECT_ID` in .env)"
-            )
+            raise Exception("PROJECT_ID must be set (argument or `GOOGLE_PROJECT_ID` in .env)")
         if not self.model_name:
             raise Exception("MODEL must be set (argument or `MODEL` in .env)")
         if not self.location:
-            raise Exception(
-                "LOCATION must be set (argument or `VERTEX_LOCATION` in .env)"
-            )
+            raise Exception("LOCATION must be set (argument or `VERTEX_LOCATION` in .env)")
 
         vertexai.init(
             project=self.project_id,
             location=self.location,
-            credentials=self.credentials,
+            credentials=self.credentials
         )
 
         self.tolerance = tolerance

tests/models/test_Gemini.py

@@ -53,18 +53,18 @@ class TestGeminiDepModelConfig(unittest.TestCase):
     def test_gemini_deployed_model_config(self):
         "Test if model config parameters passed as arguments are used"
         model = GeminiModel(
-            project_id="test_project_id", api_key="secret-key", model_name="gemini-test"
+            project_id="test_project_id", model_name="gemini-test", location="us-central1"
         )
 
         self.assertEqual(model.project_id, "test_project_id")
-        self.assertEqual(model.api_key, "secret-key")
+        self.assertEqual(model.location, "us-central1")
         self.assertEqual(model.model_name, "gemini-test")
 
     @patch.dict(
         "os.environ",
         {
             "GOOGLE_PROJECT_ID": "test_project_id",
-            "GOOGLE_API_KEY": "secret-key",
+            "LOCATION": "us-central1",
             "MODEL": "gemini-test",
         },
     )
@@ -73,23 +73,23 @@ def test_gemini_deployed_model_config_env_var(self):
         model = GeminiModel()
 
         self.assertEqual(model.project_id, "test_project_id")
-        self.assertEqual(model.api_key, "secret-key")
+        self.assertEqual(model.location, "us-central1")
         self.assertEqual(model.model_name, "gemini-test")
 
     @patch.dict(
         "os.environ",
         {
             "GOOGLE_PROJECT_ID": "test_project_id",
-            "GOOGLE_API_KEY": "secret-env-key",
+            "LOCATION": "us-central1",
             "MODEL": "gemini-test",
         },
     )
     def test_gemini_deployed_model_config_priority(self):
         "Test if model config parameters passed directly get priority"
         model = GeminiModel(
-            project_id="test_project_id", api_key="secret-key", model_name="gemini_test"
+            project_id="test_project_id", model_name="gemini_test", location="us-central1"
         )
 
-        self.assertEqual(model.project_id, "test_project_id")
-        self.assertEqual(model.api_key, "secret-key")
+        self.assertEqual(model.project_id, "test_project_id")        
+        self.assertEqual(model.location, "us-central1")
         self.assertEqual(model.model_name, "gemini_test")