From 84a540f2d6a4c41b8d06887094446417e2355f85 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 21 Mar 2024 08:27:23 -0700
Subject: [PATCH] build: fix mypy build issues

---
 .pre-commit-config.yaml                  | 16 ++++++++--------
 litellm/llms/gemini.py                   |  4 ++--
 litellm/llms/palm.py                     |  2 +-
 litellm/llms/prompt_templates/factory.py |  2 +-
 litellm/llms/vertex_ai.py                | 12 ++++++------
 5 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 44ffa8b5371..2a84048e0b1 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -16,11 +16,11 @@ repos:
         name: Check if files match
         entry: python3 ci_cd/check_files_match.py
         language: system
-# -   repo: local
-#     hooks:
-#     -   id: mypy
-#         name: mypy
-#         entry: python3 -m mypy --ignore-missing-imports
-#         language: system
-#         types: [python]
-#         files: ^litellm/
\ No newline at end of file
+-   repo: local
+    hooks:
+    -   id: mypy
+        name: mypy
+        entry: python3 -m mypy --ignore-missing-imports
+        language: system
+        types: [python]
+        files: ^litellm/
\ No newline at end of file
diff --git a/litellm/llms/gemini.py b/litellm/llms/gemini.py
index 8876c49bf20..126569ecc0c 100644
--- a/litellm/llms/gemini.py
+++ b/litellm/llms/gemini.py
@@ -118,7 +118,7 @@ def completion(
     logger_fn=None,
 ):
     try:
-        import google.generativeai as genai
+        import google.generativeai as genai  # type: ignore
     except:
         raise Exception(
             "Importing google.generativeai failed, please run 'pip install -q google-generativeai"
@@ -308,7 +308,7 @@ async def async_completion(
     messages,
     encoding,
 ):
-    import google.generativeai as genai
+    import google.generativeai as genai  # type: ignore
 
     response = await _model.generate_content_async(
         contents=prompt,
diff --git a/litellm/llms/palm.py b/litellm/llms/palm.py
index 060e6dca150..3f0664b1d54 100644
--- a/litellm/llms/palm.py
+++ b/litellm/llms/palm.py
@@ -98,7 +98,7 @@ def completion(
     logger_fn=None,
 ):
     try:
-        import google.generativeai as palm
+        import google.generativeai as palm  # type: ignore
     except:
         raise Exception(
             "Importing google.generativeai failed, please run 'pip install -q google-generativeai"
diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py
index 87c6d89610f..62a387e0b39 100644
--- a/litellm/llms/prompt_templates/factory.py
+++ b/litellm/llms/prompt_templates/factory.py
@@ -923,7 +923,7 @@ def gemini_text_image_pt(messages: list):
     }
     """
     try:
-        import google.generativeai as genai
+        import google.generativeai as genai  # type: ignore
     except:
         raise Exception(
             "Importing google.generativeai failed, please run 'pip install -q google-generativeai"
diff --git a/litellm/llms/vertex_ai.py b/litellm/llms/vertex_ai.py
index 0a7980fda2f..a38530f151d 100644
--- a/litellm/llms/vertex_ai.py
+++ b/litellm/llms/vertex_ai.py
@@ -289,11 +289,11 @@ def completion(
             Part,
             GenerationConfig,
         )
-        from google.cloud import aiplatform
+        from google.cloud import aiplatform  # type: ignore
         from google.protobuf import json_format  # type: ignore
         from google.protobuf.struct_pb2 import Value  # type: ignore
-        from google.cloud.aiplatform_v1beta1.types import content as gapic_content_types
-        import google.auth
+        from google.cloud.aiplatform_v1beta1.types import content as gapic_content_types  # type: ignore
+        import google.auth  # type: ignore
 
         ## Load credentials with the correct quota project ref: https://github.com/googleapis/python-aiplatform/issues/2557#issuecomment-1709284744
         print_verbose(
@@ -783,7 +783,7 @@ async def async_completion(
             """
             Vertex AI Model Garden
             """
-            from google.cloud import aiplatform
+            from google.cloud import aiplatform  # type: ignore
 
             ## LOGGING
             logging_obj.pre_call(
@@ -969,7 +969,7 @@ async def async_streaming(
         )
         response = llm_model.predict_streaming_async(prompt, **optional_params)
     elif mode == "custom":
-        from google.cloud import aiplatform
+        from google.cloud import aiplatform  # type: ignore
 
         stream = optional_params.pop("stream", None)
 
@@ -1059,7 +1059,7 @@ def embedding(
         )
 
     from vertexai.language_models import TextEmbeddingModel
-    import google.auth
+    import google.auth  # type: ignore
 
     ## Load credentials with the correct quota project ref: https://github.com/googleapis/python-aiplatform/issues/2557#issuecomment-1709284744
     try: