From 3f348fa38b20213ff7cbe95da11633cb0fa0dc74 Mon Sep 17 00:00:00 2001 From: Mark Daoust Date: Tue, 24 Sep 2024 09:28:14 -0700 Subject: [PATCH] fix typing Change-Id: If892b20ca29d1afb82c48ae1a49bef58e0421bab --- google/generativeai/types/content_types.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/google/generativeai/types/content_types.py b/google/generativeai/types/content_types.py index 37e46fd14..39f0e138a 100644 --- a/google/generativeai/types/content_types.py +++ b/google/generativeai/types/content_types.py @@ -671,7 +671,16 @@ 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): @@ -679,7 +688,7 @@ def _make_google_search_retrieval(gsr: GoogleSearchRetrievalType): 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)