google-gemini · MarkDaoust · Dec 20, 2023 · Dec 12, 2023 · Dec 20, 2023 · Dec 20, 2023
diff --git a/google/generativeai/types/citation_types.py b/google/generativeai/types/citation_types.py
@@ -13,13 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from __future__ import annotations
+from typing import List
 
-from typing import Optional, List
+from typing_extensions import TypedDict
 
 from google.ai import generativelanguage as glm
 from google.generativeai import string_utils
 
-from typing import TypedDict
 
 __all__ = [
     "CitationMetadataDict",

diff --git a/google/generativeai/types/discuss_types.py b/google/generativeai/types/discuss_types.py
@@ -16,14 +16,16 @@
 
 import abc
 import dataclasses
-from typing import Any, Dict, TypedDict, Union, Iterable, Optional, Tuple, List
+from typing import Any, Dict, Union, Iterable, Optional, Tuple, List
+from typing_extensions import TypedDict
 
 import google.ai.generativelanguage as glm
 from google.generativeai import string_utils
 
 from google.generativeai.types import safety_types
 from google.generativeai.types import citation_types
 
+
 __all__ = [
     "MessageDict",
     "MessageOptions",

diff --git a/google/generativeai/types/model_types.py b/google/generativeai/types/model_types.py
@@ -22,12 +22,16 @@
 import json
 import pathlib
 import re
-from typing import Any, Iterable, TypedDict, Union
+
+from typing import Any, Iterable, Union
+
 import urllib.request
+from typing_extensions import TypedDict
 
 import google.ai.generativelanguage as glm
 from google.generativeai import string_utils
 
+
 __all__ = [
     "Model",
     "ModelNameOptions",
@@ -266,12 +270,12 @@ def _convert_dict(data, input_key, output_key):
 
     try:
         inputs = data[input_key]
-    except KeyError as e:
+    except KeyError:
         raise KeyError(f'input_key is "{input_key}", but data has keys: {sorted(data.keys())}')
 
     try:
         outputs = data[output_key]
-    except KeyError as e:
+    except KeyError:
         raise KeyError(f'output_key is "{output_key}", but data has keys: {sorted(data.keys())}')
 
     for i, o in zip(inputs, outputs):

diff --git a/google/generativeai/types/safety_types.py b/google/generativeai/types/safety_types.py
@@ -16,11 +16,15 @@
 
 from collections.abc import Mapping
 
+import typing
+from typing import Dict, Iterable, List, Union
+
+from typing_extensions import TypedDict
+
+
 from google.ai import generativelanguage as glm
 from google.generativeai import string_utils
 
-import typing
-from typing import Iterable, Dict, Iterable, List, TypedDict, Union
 
 __all__ = [
     "HarmCategory",

diff --git a/google/generativeai/types/text_types.py b/google/generativeai/types/text_types.py
@@ -14,9 +14,11 @@
 # limitations under the License.
 from __future__ import annotations
 
+import sys
 import abc
 import dataclasses
-from typing import Any, Dict, List, TypedDict
+from typing import Any, Dict, List
+from typing_extensions import TypedDict
 
 from google.generativeai import string_utils
 from google.generativeai.types import safety_types

diff --git a/setup.py b/setup.py
@@ -46,6 +46,7 @@ def get_version():
     "google-ai-generativelanguage==0.3.3",
     "google-auth",
     "google-api-core",
+    "typing-extensions",
     "protobuf",
     "tqdm",
 ]