feat(py): add Gemini 3.1 preview support to python #5032
feat(py): add Gemini 3.1 preview support to python #5032cabljac wants to merge 2 commits intogenkit-ai:mainfrom
Conversation
Construct Google GenAI clients with explicit typed keyword arguments instead of unpacked kwargs maps, and apply Ruff formatting to plugin tests. Made-with: Cursor
github-actions
bot
added
docs
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for Gemini 3.1 preview models (Flash Lite, Pro, and Pro with Custom Tools) across the Google AI and Vertex AI plugins. A key architectural change includes a routing mechanism in the Vertex AI plugin to ensure Gemini 3.1 preview models use the 'global' location, while maintaining configured locations for other models. The PR also updates documentation, enums, and samples to reflect these additions. Feedback was provided regarding the newly defined ModelInfo constants, which are currently not registered in the SUPPORTED_MODELS dictionary, potentially causing them to fall back to generic metadata in the UI.
| GEMINI_3_1_FLASH_LITE_PREVIEW = ModelInfo( | ||
| label='Google AI - Gemini 3.1 Flash Lite Preview', | ||
| supports=Supports( | ||
| multiturn=True, | ||
| media=True, | ||
| tools=True, | ||
| tool_choice=True, | ||
| system_role=True, | ||
| constrained=Constrained.NO_TOOLS, | ||
| output=['text', 'json'], | ||
| ), | ||
| ) | ||
|
|
||
| GEMINI_3_1_PRO_PREVIEW_CUSTOMTOOLS = ModelInfo( | ||
| label='Google AI - Gemini 3.1 Pro Preview Custom Tools', | ||
| supports=Supports( | ||
| multiturn=True, | ||
| media=True, | ||
| tools=True, | ||
| tool_choice=True, | ||
| system_role=True, | ||
| constrained=Constrained.NO_TOOLS, | ||
| output=['text', 'json'], | ||
| ), | ||
| ) | ||
|
|
||
| GEMINI_3_1_PRO_PREVIEW = ModelInfo( | ||
| label='Google AI - Gemini 3.1 Pro Preview', | ||
| supports=Supports( | ||
| multiturn=True, | ||
| media=True, | ||
| tools=True, | ||
| tool_choice=True, | ||
| system_role=True, | ||
| constrained=Constrained.NO_TOOLS, | ||
| output=['text', 'json'], | ||
| ), | ||
| ) |
There was a problem hiding this comment.
While these ModelInfo constants are correctly defined, they are currently not being used for model lookup because they aren't added to the SUPPORTED_MODELS dictionary. This means that when these models are resolved, they will fall back to a generic ModelInfo with a default label and support flags, rather than using the specific metadata defined here.
Consider populating the SUPPORTED_MODELS dictionary with these constants (and existing ones) to ensure the correct metadata is surfaced in the Genkit UI and used during execution.
cabljac
changed the title
cabljac
changed the title
cabljac
changed the title
huangjeff5
changed the title
| ai = Genkit( | ||
| plugins=[GoogleAI()], | ||
| model=f'googleai/{gemini.GoogleAIGeminiVersion.GEMINI_3_FLASH_PREVIEW}', | ||
| model=f'googleai/{gemini.GoogleAIGeminiVersion.GEMINI_3_1_PRO_PREVIEW_CUSTOMTOOLS}', |
There was a problem hiding this comment.
Do we need it to be the CUSTOMTOOLS preview? Is tool calling not available in the regular 3.1 pro model?
| try: | ||
| response = await draw_image_with_imagen() | ||
| print(response.model_dump_json(indent=2)) # noqa: T201 | ||
| print(await summarize_with_gemini_31_pro()) # noqa: T201 |
There was a problem hiding this comment.
Think we can skip adding these flows in the vertexai sample since it seems to be showing off basic generate usage, and this is the Imagen sample.
In general I think updating existing samples to use the latest model is good enough, unless there are some major new features to demo.
|
Do the new models show up correctly in dev UI? |
|
The two-client approach is going to get messy since we have to duplicate all those arguments. I'd suggest a location-keyed client cache instead: # __init__: store typed args, drop _global_runtime_client
self._api_key = api_key
self._credentials = credentials
self._debug_config = debug_config
self._http_options = _inject_attribution_headers(http_options, base_url, api_version)
self._client_by_location: dict[str, Callable[[], genai.client.Client]] = {}
self._runtime_client = self._client_getter(self._location)
def _make_client(self, location: str) -> genai.client.Client:
return genai.client.Client(
vertexai=True,
api_key=self._api_key,
credentials=self._credentials,
project=self._project,
location=location,
debug_config=self._debug_config,
http_options=self._http_options,
)
def _client_getter(self, location: str) -> Callable[[], genai.client.Client]:
if location not in self._client_by_location:
self._client_by_location[location] = loop_local_client(
lambda loc=location: self._make_client(loc)
)
return self._client_by_location[location]Rename def _preferred_vertex_location(name: str, default: str) -> str:
lower = name.lower()
if lower.startswith('gemini-3.1-') and 'preview' in lower:
return 'global'
return defaultAnd there's a subtle bug in the current async def _run(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse:
client = self._client_getter(_preferred_vertex_location(clean_name, self._location))()
if clean_name.lower().startswith('image'):
model = ImagenModel(clean_name, client)
elif is_veo_model(clean_name):
model = VeoModel(clean_name, client)
else:
model = GeminiModel(clean_name, client)
return await model.generate(request, ctx) |
4 participants
Summary
Adds Python support for new Gemini 3.1 preview models:
Updates Vertex plugin routing so Gemini 3.1 preview models default to global location (while non-3.1-preview models keep configured region).
Updates samples/Dev UI flows so all 5 models can be exercised directly:
Adds/updates tests for enum support and Vertex global-routing behavior.
Issue
Checklist: