diff --git a/google/generativeai/types/content_types.py b/google/generativeai/types/content_types.py index af6ebafad..f82e26839 100644 --- a/google/generativeai/types/content_types.py +++ b/google/generativeai/types/content_types.py @@ -703,11 +703,11 @@ def __init__( self, *, function_declarations: Iterable[FunctionDeclarationType] | None = None, - google_search_retrieval: Union[protos.GoogleSearchRetrieval, str] | None = None, + google_search_retrieval: GoogleSearchRetrievalType | None = None, code_execution: protos.CodeExecution | None = None, ): # The main path doesn't use this but is seems useful. - if function_declarations: + if function_declarations is not None: self._function_declarations = [ _make_function_declaration(f) for f in function_declarations ] @@ -722,15 +722,10 @@ def __init__( self._function_declarations = [] self._index = {} - if google_search_retrieval: - if isinstance(google_search_retrieval, str): - self._google_search_retrieval = { - "google_search_retrieval": { - "dynamic_retrieval_config": {"mode": to_mode(google_search_retrieval)} - } - } - else: - self._google_search_retrieval = _make_google_search_retrieval(google_search_retrieval) + if google_search_retrieval is not None: + self._google_search_retrieval = _make_google_search_retrieval(google_search_retrieval) + else: + self._google_search_retrieval = None self._proto = protos.Tool( function_declarations=[_encode_fd(fd) for fd in self._function_declarations], diff --git a/setup.py b/setup.py index 29841ba1d..0575dcd28 100644 --- a/setup.py +++ b/setup.py @@ -42,7 +42,7 @@ def get_version(): release_status = "Development Status :: 5 - Production/Stable" dependencies = [ - "google-ai-generativelanguage==0.6.9", + "google-ai-generativelanguage@https://storage.googleapis.com/generativeai-downloads/preview/ai-generativelanguage-v1beta-py.tar.gz", "google-api-core", "google-api-python-client", "google-auth>=2.15.0", # 2.15 adds API key auth support