Update Gemini with new SDK
Purpose
Gemini have introduced a new Python SDK. This PR is to migrate from the deprecated google-generativeai package to the newer google-genai package.
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? - [x] Update Auth https://ai.google.dev/gemini-api/docs/migrate#authenticate
- [x] async - handled differently in new SDK https://ai.google.dev/gemini-api/docs/migrate#async Check our implementation still makes sense.
- [x] Replace
GenerativeModelwith a client instance, created once, earlier in the pipeline (including updating tests as required). - [x] In
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. - [x] For accuracy of the comments in the test
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). - [x] The existing code supports "Per model API keys". This seems to go against the grain of how the Client authentication works - should this be maintained?
- Probably not required, but easy enough to maintain, so I've left it in for now.
- [x] SafetySetting (https://googleapis.github.io/python-genai/genai.html#genai.types.SafetySetting), includes a
methodparameter. We don't appear to be using this, or any equivalent, in the current code. Do we need to add this?- Not required
- [x] Fix test
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.
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.
Codecov Report
Attention: Patch coverage is 92.70833% with 7 lines in your changes missing coverage. Please review.
Project coverage is 67.24%. Comparing base (
227ed0d) to head (6c975bc).
| Files with missing lines | Patch % | Lines |
|---|---|---|
| src/prompto/upload_media.py | 56.25% | 7 Missing :warning: |
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
:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.
:rocket: New features to boost your workflow:
- :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.