prompto icon indicating copy to clipboard operation
prompto copied to clipboard

Update Gemini with new SDK

Open andrewphilipsmith opened this issue 7 months ago • 2 comments

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)

  • [ ] vertexai is dependent on the deprecated google-generativeai package. 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 GenerativeModel with a client instance, created once, earlier in the pipeline (including updating tests as required).
  • [x] In test_gemini.py test_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_inputs to 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 of google.genai.client.aio.models.generate_content and google.genai.client.models.generate_content (both return a genai.types.GenerateContentResponse object).
  • [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 method parameter. 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_var This 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.

andrewphilipsmith avatar Apr 23 '25 15:04 andrewphilipsmith

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.

andrewphilipsmith avatar Apr 23 '25 15:04 andrewphilipsmith

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.

codecov-commenter avatar Apr 29 '25 16:04 codecov-commenter