fix typing

Change-Id: If892b20ca29d1afb82c48ae1a49bef58e0421bab
google-gemini · Sep 24, 2024 · 3f348fa · 3f348fa
1 parent af6ce40
commit 3f348fa
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/google/generativeai/types/content_types.py b/google/generativeai/types/content_types.py
@@ -671,15 +671,24 @@ def _encode_fd(fd: FunctionDeclaration | protos.FunctionDeclaration) -> protos.F
     return fd.to_proto()
 
 
-GoogleSearchRetrievalType = Union[protos.GoogleSearchRetrieval, dict[str, float]]
+class DynamicRetrievalConfigDict(TypedDict):
+    mode: protos.DynamicRetrievalConfig.mode
+    dynamic_threshold: float
+
+DynamicRetrievalConfig = Union[protos.DynamicRetrievalConfig, DynamicRetrievalConfigDict]
+
+class GoogleSearchRetrievalDict(TypedDict):
+    dynamic_retrieval_config: DynamicRetrievalConfig
+
+GoogleSearchRetrievalType = Union[protos.GoogleSearchRetrieval, GoogleSearchRetrievalDict]
 
 
 def _make_google_search_retrieval(gsr: GoogleSearchRetrievalType):
     if isinstance(gsr, protos.GoogleSearchRetrieval):
         return gsr
     elif isinstance(gsr, Mapping):
         drc = gsr.get("dynamic_retrieval_config", None)
-        if drc is not None:
+        if drc is not None and isinstance(drc, Mapping):
             mode = drc.get("mode", None)
             if mode is not None:
                 mode = to_mode(mode)