-
Notifications
You must be signed in to change notification settings - Fork 2.8k
feat: add custom client support for Gemini #4345
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?
feat: add custom client support for Gemini #4345
Conversation
Summary of ChangesHello @challenger71498, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
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.
Code Review
This pull request introduces a valuable feature by adding support for custom google.genai.Client instances to the Gemini model class. This greatly improves flexibility for testing, custom authentication, and other advanced configurations. The implementation is clean, and the accompanying unit tests are thorough and well-written. I have a couple of minor suggestions to improve the code samples in the docstrings to make them more complete for users.
Please ensure you have read the contribution guide before creating a pull request.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
2. Or, if no issue exists, describe the change:
Problem:
Users need the ability to inject custom Google GenAI Client instances for testing, custom authentication, or different configurations. Currently, the
Geminimodel class always creates its own internal clients for both standard API calls and Live API streaming, making it difficult to override client behavior without modifying internal code.Solution:
Added two new optional parameters to the
Geminimodel class:custom_api_client: Allows injecting a custom client forgenerate_contentoperationscustom_live_api_client: Allows injecting a custom client for Live API (bi-directional streaming) operationsWhen these parameters are provided, they override the default client creation logic in the
api_clientand_live_api_clientproperties respectively. This enables users to:Testing Plan
Unit Tests:
Summary of pytest results:
Added two comprehensive unit tests in
tests/unittests/models/test_google_llm.py:test_custom_api_client_is_used_for_generate_content- Verifies thatcustom_api_clientis properly used for standard API operationstest_custom_live_api_client_is_used_for_connect- Verifies thatcustom_live_api_clientis properly used for Live API streaming operationsBoth tests use mock clients to verify the injection mechanism works correctly without making actual API calls.
Manual End-to-End (E2E) Tests:
To manually test these changes:
Checklist
Additional context
Implementation details:
Field(exclude=True, frozen=True, repr=False)to ensure they don't interfere with Pydantic serialization, are immutable after creation, and don't clutter string representationsOptional[Client]and default toNoneNone, the existing client creation logic is used (backward compatible)Clientimport fromTYPE_CHECKINGblock to top-level imports since it's now used in type annotations for runtime fields