diff --git a/setup.py b/setup.py index 0575dcd28..89af61515 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@https://storage.googleapis.com/generativeai-downloads/preview/ai-generativelanguage-v1beta-py.tar.gz", + "google-ai-generativelanguage==0.6.5", "google-api-core", "google-api-python-client", "google-auth>=2.15.0", # 2.15 adds API key auth support diff --git a/tests/test_generative_models.py b/tests/test_generative_models.py index 73789346d..c4d46ffec 100644 --- a/tests/test_generative_models.py +++ b/tests/test_generative_models.py @@ -842,7 +842,10 @@ def test_count_tokens_smoke(self, kwargs): self.responses["count_tokens"] = [protos.CountTokensResponse(total_tokens=7)] model = generative_models.GenerativeModel("gemini-pro-vision", system_instruction=si) response = model.count_tokens(**kwargs) - self.assertEqual(type(response).to_dict(response), {"total_tokens": 7}) + self.assertEqual( + type(response).to_dict(response, including_default_value_fields=False), + {"total_tokens": 7}, + ) @parameterized.named_parameters( [ diff --git a/tests/test_generative_models_async.py b/tests/test_generative_models_async.py index 03055ffb3..3dcf49ae4 100644 --- a/tests/test_generative_models_async.py +++ b/tests/test_generative_models_async.py @@ -216,7 +216,10 @@ async def test_count_tokens_smoke(self, contents): self.responses["count_tokens"] = [protos.CountTokensResponse(total_tokens=7)] model = generative_models.GenerativeModel("gemini-pro-vision") response = await model.count_tokens_async(contents) - self.assertEqual(type(response).to_dict(response), {"total_tokens": 7}) + self.assertEqual( + type(response).to_dict(response, including_default_value_fields=False), + {"total_tokens": 7}, + ) async def test_stream_generate_content_called_with_request_options(self): self.client.stream_generate_content = unittest.mock.AsyncMock()