Generator: Add support for Google Gemini models #1306
Conversation
There was a problem hiding this comment.
Working with GCP based services has shown that auth via api_key is not always as straight forward as providing a single environment variable value. Can you provide details on how the key used would be scoped or generated.
Also consider enabling the google library to attempt auth even when no key is provided.
garak/generators/gemini.py
Outdated
| responses = [] | ||
| import logging | ||
|
|
||
| for _ in range(generations_this_call): |
There was a problem hiding this comment.
Iterating generations like this is will not combine with backoff correctly. As written any raised backoff exception will throw away all completed generations and start over.
Looking at gemini docs multiple generations can be obtained in a single call using GenerateContentConfig by setting candidateCount= generations_this_call and passing generate_content() a config named parameter.
If calling for more than one generation please validate how the response object will be formatted.
There was a problem hiding this comment.
Addressed - The implementation now uses candidate_count=generations_this_call in GenerateContentConfig and makes a single API call via _generate_content_with_backoff(). This ensures backoff retries don't discard completed generations. The response format has been validated and properly handles multiple candidates from the API response.
1. Update Gemini generator to handle test model names gracefully 2. Add missing documentation file for Gemini generator 3. Add Gemini generator to documentation toctree 4. Add google-generativeai dependency to pyproject.toml
remove api validation Co-authored-by: Jeffrey Martin <[email protected]> Signed-off-by: Divya Chitimalla <[email protected]>
- Update _call_model to use Gemini's native candidateCount parameter for multiple generations - Process response candidates correctly to extract text from each generation - Remove generation config from model initialization and set it per request - Fix backoff handling to properly retry the entire batch of generations - Ensure consistent number of responses are returned
There was a problem hiding this comment.
General comments, the auth concern may be reasonable to defer to require environment variable based API key for initial acceptance.
Another concern that may be important to address:
This PR uses the package google-generativeai however current docs recommend using google-genai which has the same imported package name but a different interface.
tests/generators/test_gemini.py
Outdated
|
|
||
| # Create the generator with a native audio model | ||
| generator = GeminiGenerator(name="gemini-2.5-flash-native-audio") | ||
| output = generator._call_model("Transcribe this text.") |
There was a problem hiding this comment.
This does not make sense the input modality suggested by this test is audio not text.
The current generator also inherits the default modality, {"in": {"text"}, "out": {"text"}} unless overridden this generator should only accept text prompts.
Co-authored-by: Jeffrey Martin <[email protected]> Signed-off-by: Divya Chitimalla <[email protected]>
Co-authored-by: Jeffrey Martin <[email protected]> Signed-off-by: Divya Chitimalla <[email protected]>
Co-authored-by: Jeffrey Martin <[email protected]> Signed-off-by: Divya Chitimalla <[email protected]>
Co-authored-by: Jeffrey Martin <[email protected]> Signed-off-by: Divya Chitimalla <[email protected]>
Fixed Failing Tests
Migrating it to the new google-genai , Support both API Keys and Vertex AI (ADC) and Improving the backoff logic
|
DCO Assistant Lite bot All contributors have signed the DCO ✍️ ✅ |
|
I have read the DCO Document and I hereby sign the DCO |
|
recheck |
|
Faced issues trying to run gemini models using Litellm. This generator natively supports google models using Google’s official google.generativeai library—more reliable than LiteLLM. It supports multiple models (2.5 Pro, Flash, etc.) with error handling and model-specific config options. #443