fix(py): use async Google GenAI model discovery - #6298
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces asynchronous model discovery to prevent blocking the event loop when fetching models from the Google GenAI SDK. This is achieved by wrapping the synchronous model listing call in asyncio.to_thread and updating the plugin initialization and action listing methods to await this new asynchronous helper. Feedback was provided to refactor the newly added test to avoid potential flakiness in CI/CD environments by asserting thread execution context instead of relying on timing delays.
|
Two issues found, One blocking and one nit. Blocking: five tests in test/google_plugin_test.py fail at this head with "object MagicMock can't be used in 'await' expression" (test_googleai_initialize, test_googleai_list_actions, test_vertexai_initialize, test_vertexai_list_actions, test_vertexai_list_actions_without_supported_actions). They stub client.models.list but init() and list_actions() now go through client.aio.models.list. Can we point those stubs at the aio surface, or add a fixture that does it once? Nit: the PR body still describes the worker-thread approach and the slow-call test. Worth updating to match the client.aio change. |
Dropping the _list_known_* tests took the only test that drove a generateVideos model through the Google AI categorization branch. The Vertex test returns before supported_actions is read, so it cannot reach it.
Summary
client.aio.models.list()for Google GenAI model discovery in asynchronous plugin entry points.Fixes #6186
Tests
pytest -p no:cacheprovider -o addopts='' py/packages/genkit-google-genai/test py/packages/genkit-google-genai/tests -q— 738 passedruff checkon the changed Google GenAI source and test files — passedpython -m compileallon the Google GenAI source and test packages — passedgit diff --checkon the changed files — passedCompatibility / Known limitations
client.aio.models.list()for both Google AI and Vertex AI plugin entry points.