fix tests

Change-Id: Ib9e19d78861da180f713e09ec93d366d5d7b5762
google-gemini · Sep 20, 2024 · e60841e · e60841e
1 parent c5cebf2
commit e60841e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 12 deletions.
diff --git 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
@@ -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