-
Notifications
You must be signed in to change notification settings - Fork 2
Update Gemini with new SDK #121
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
|
The first commit updates the relevant tests to use the new SDK, without updating the actual implementation. There will be many failing tests for now. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull Request Overview
This PR migrates the Gemini API integration from the deprecated google-generativeai package to the newer google-genai package by updating client instantiation, safety settings handling, and test expectations.
- Updated tests and patches to reflect the new SDK classes and method signatures.
- Refactored GeminiAPI to use Client and GenerateContentConfig from google-genai and updated dependency management in pyproject.toml.
Reviewed Changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| tests/apis/gemini/*.py | Updated test patches and assertions to reflect new SDK changes |
| src/prompto/apis/gemini/gemini.py | Refactored client instantiation and method interfaces for the new SDK |
| pyproject.toml | Updated dependency versions and deprecated package removal notices |
src/prompto/apis/gemini/gemini.py
Outdated
| HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE, | ||
| HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE, | ||
| } | ||
| SafetySetting() |
Copilot
AI
Apr 24, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The standalone call to SafetySetting() in the 'none' safety_filter branch is unnecessary and may confuse readers. Please remove this extraneous call.
| SafetySetting() |
src/prompto/apis/gemini/gemini.py
Outdated
| """ | ||
| if prompt_dict["prompt"][0]["role"] == "system": | ||
| prompt, model_name, model, safety_settings, generation_config = ( | ||
| prompt, model_name, client, safety_settings, generation_config = ( |
Copilot
AI
Apr 24, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
In _query_history, the tuple unpacking does not match the updated return type from _obtain_model_inputs (Client and GenerateContentConfig). Update the unpacking (e.g. unpack as (prompt, model_name, client, generation_config, _)) and remove or adjust the usage of the safety_settings parameter in the subsequent send_message_async call.
| prompt, model_name, client, safety_settings, generation_config = ( | |
| prompt, model_name, client, generation_config, _ = ( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull Request Overview
This PR migrates the Gemini integration from the deprecated google-generativeai package to the new google-genai SDK, updating imports, client initialization, safety settings, and related test cases. Key changes include:
- Updating all API calls and client interactions to use google-genai (e.g., AsyncClient, GenerateContentConfig) instead of the deprecated package.
- Adjusting test mocks and expected responses for the new SDK APIs and safety settings.
- Modifying utility functions (e.g., parsing multimedia parts and converting history dictionaries) to accept the new client parameter and structure.
Reviewed Changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| tests/apis/gemini/*.py | Updated tests with new client imports and proper mocks for async calls |
| src/prompto/apis/gemini/gemini_utils.py | Modified utility functions signature to include a client parameter |
| src/prompto/apis/gemini/gemini.py | Migrated API client usage and safety settings configuration to new SDK |
| pyproject.toml | Updated dependency from google-generativeai to google-genai |
| if isinstance(part, str): | ||
| return part | ||
| # return part | ||
| print(f"Part is a string: {part}") |
Copilot
AI
Apr 29, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Remove the debug print statement used when handling string parts, as it may leak unwanted output in production.
| print(f"Part is a string: {part}") |
src/prompto/apis/gemini/gemini.py
Outdated
| # No need to send the generation_config again, as it is no different | ||
| # from the one used to create the chat | ||
| last_msg = prompt[-1] | ||
| print(f"whole prompt: {prompt}") |
Copilot
AI
Apr 29, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Remove or replace debug print statements in the _query_history method before production, to ensure cleaner logging and performance.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #121 +/- ##
==========================================
+ Coverage 65.09% 67.24% +2.15%
==========================================
Files 44 44
Lines 2679 2690 +11
==========================================
+ Hits 1744 1809 +65
+ Misses 935 881 -54 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull Request Overview
This PR migrates the Gemini API integration from the deprecated google-generativeai package to the new google-genai package. Key changes include updating API call methods to use AsyncClient and related async methods, revising safety settings to use new types (SafetySetting) and updating tests and dependency definitions.
Reviewed Changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| tests/apis/gemini/*.py | Updated tests to patch new client methods and use updated types, ensuring compatibility with the new SDK. |
| src/prompto/apis/gemini/*.py | Replaced usage of GenerativeModel with Client instances and updated safety settings and API calls accordingly. |
| src/prompto/upload_media.py | Updated media upload functions to use new asynchronous SDK methods and adjusted settings creation. |
| pyproject.toml | Modified dependency versions and extras to reflect changes in packages and SDK versions. |
| For now, we just create a temporary directory for the data folder. | ||
| """ | ||
| # TODO: A better solution would be to create an option in the | ||
| # Settings constructor to not create the directories. | ||
| # But for now we'll just pass it a temporary directory. | ||
| import tempfile | ||
|
|
||
| with tempfile.TemporaryDirectory() as temp_dir: | ||
| data_folder = os.path.join(temp_dir, "data") | ||
| os.makedirs(data_folder, exist_ok=True) | ||
| dummy_settings = Settings(data_folder=data_folder) |
Copilot
AI
May 6, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Creating a Settings instance inside a TemporaryDirectory context may lead to issues since the temporary directory will be deleted after the with block exits. Consider creating a persistent directory or allowing Settings to manage its own directory lifecycle.
| For now, we just create a temporary directory for the data folder. | |
| """ | |
| # TODO: A better solution would be to create an option in the | |
| # Settings constructor to not create the directories. | |
| # But for now we'll just pass it a temporary directory. | |
| import tempfile | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| data_folder = os.path.join(temp_dir, "data") | |
| os.makedirs(data_folder, exist_ok=True) | |
| dummy_settings = Settings(data_folder=data_folder) | |
| A persistent directory is created for the data folder. | |
| """ | |
| # Define a persistent directory for the data folder | |
| data_folder = os.path.expanduser("~/.prompto/data") | |
| os.makedirs(data_folder, exist_ok=True) | |
| # Create the Settings instance with the persistent data folder | |
| dummy_settings = Settings(data_folder=data_folder) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I have not fixed this in this pull request, but created a separate issue to track it: #122
Purpose
Gemini have introduced a new Python SDK. This PR is to migrate from the deprecated
google-generativeaipackage to the newergoogle-genaipackage.See references here https://ai.google.dev/api?lang=python
Migration Guide https://ai.google.dev/gemini-api/docs/migrate
Notes / Questions / Todos:
(In no particular order)
vertexaiis dependent on the deprecatedgoogle-generativeaipackage. Do we need to upgrade this too?GenerativeModelwith a client instance, created once, earlier in the pipeline (including updating tests as required).test_gemini.pytest_gemini_obtain_model_inputs(): Have updated the tests to assume that the methods will return a slightly different tuple of values. For now, assume that the most sensible thing for the_obtain_model_inputsto return here is the AsyncClient instance. It may be that returning nothing is the sensible thing to do, in which case, different parts of the test will need to be updated.test_gemini_query_string, check if there is a difference in the return type ofgoogle.genai.client.aio.models.generate_contentandgoogle.genai.client.models.generate_content(both return agenai.types.GenerateContentResponseobject).methodparameter. We don't appear to be using this, or any equivalent, in the current code. Do we need to add this?tests/apis/gemini/test_gemini_chat_input.py::test_gemini_query_chat_no_env_varThis test passes when executed alone, but fails when executed with all tests. This is probably due to the environment variable being monkeypatched somewhere without being reset / properly scoped.