From 04dbd5fd615fd6e5a35b070eac4644cec74b23e6 Mon Sep 17 00:00:00 2001
From: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com>
Date: Wed, 20 Dec 2023 01:52:25 +0100
Subject: [PATCH] Import typeddict from typing_extentions (#104)

* import typeddict from typing_extentions

* Add typing-extensions

* Move shim logic to one file.

* typo

* simpilfy

* order

---------

Co-authored-by: Mark Daoust <markdaoust@google.com>
---
 google/generativeai/types/citation_types.py |  4 ++--
 google/generativeai/types/discuss_types.py  |  4 +++-
 google/generativeai/types/model_types.py    | 10 +++++++---
 google/generativeai/types/safety_types.py   |  8 ++++++--
 google/generativeai/types/text_types.py     |  4 +++-
 setup.py                                    |  1 +
 6 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/google/generativeai/types/citation_types.py b/google/generativeai/types/citation_types.py
index e720fcd88..ae857c35b 100644
--- 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
index b6e1f6a36..0cb393e5c 100644
--- a/google/generativeai/types/discuss_types.py
+++ b/google/generativeai/types/discuss_types.py
@@ -16,7 +16,8 @@
 
 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
@@ -24,6 +25,7 @@
 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
index d628b5ef3..1aaac4beb 100644
--- 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
index 893206526..b11baec66 100644
--- a/google/generativeai/types/safety_types.py
+++ b/google/generativeai/types/safety_types.py
@@ -2,11 +2,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
index a8fa507eb..f66c0fb32 100644
--- 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
index d1568c981..a9aceeb4a 100644
--- a/setup.py
+++ b/setup.py
@@ -46,6 +46,7 @@ def get_version():
     "google-ai-generativelanguage==0.4.0",
     "google-auth",
     "google-api-core",
+    "typing-extensions",
     "protobuf",
     "tqdm",
 ]