fix(py/plugins/genai): expose full Gemini config schema (#2540)

Author: pavelgj

py/packages/genkit-ai/src/genkit/ai/__init__.py

@@ -47,6 +47,7 @@
     Role,
     Supports,
     TextPart,
+    ToolDefinition,
     ToolRequest,
     ToolRequestPart,
     ToolResponse,
@@ -85,5 +86,6 @@
     ToolResponsePart.__name__,
     GenerationCommonConfig.__name__,
     GenerationUsage.__name__,
+    ToolDefinition.__name__,
     tool_response.__name__,
 ]

py/packages/genkit-ai/src/genkit/blocks/prompt.py

@@ -339,8 +339,6 @@ def to_generate_action_options(
     model = model or registry.default_model
     if model is None:
         raise Exception('No model configured.')
-    if not isinstance(config, GenerationCommonConfig | dict | None):
-        raise AttributeError('Invalid generate config provided')
     resolved_msgs: list[Message] = []
     if system:
         resolved_msgs.append(

py/plugins/google-genai/src/genkit/plugins/google_genai/__init__.py

@@ -20,7 +20,10 @@
     GeminiEmbeddingModels,
     VertexEmbeddingModels,
 )
-from genkit.plugins.google_genai.models.gemini import GeminiVersion
+from genkit.plugins.google_genai.models.gemini import (
+    GeminiConfigSchema,
+    GeminiVersion,
+)
 
 
 def package_name() -> str:
@@ -40,4 +43,5 @@ def package_name() -> str:
     VertexEmbeddingModels.__name__,
     GeminiVersion.__name__,
     EmbeddingTaskType.__name__,
+    GeminiConfigSchema.__name__,
 ]

py/plugins/google-genai/src/genkit/plugins/google_genai/google.py

@@ -30,6 +30,7 @@
 )
 from genkit.plugins.google_genai.models.gemini import (
     GeminiApiOnlyVersion,
+    GeminiConfigSchema,
     GeminiModel,
     GeminiVersion,
 )
@@ -101,6 +102,7 @@ def initialize(self, ai: GenkitRegistry) -> None:
                 name=google_genai_name(version),
                 fn=gemini_model.generate,
                 metadata=gemini_model.metadata,
+                config_schema=GeminiConfigSchema,
             )
 
         embeding_models = (

py/plugins/google-genai/src/genkit/plugins/google_genai/models/gemini.py

@@ -124,22 +124,27 @@
 from google import genai
 from google.genai import types as genai_types
 
-from genkit.ai.registry import GenkitRegistry
-from genkit.core.action import ActionKind, ActionRunContext
-from genkit.core.typing import (
+from genkit.ai import (
+    ActionKind,
+    ActionRunContext,
     GenerateRequest,
     GenerateResponse,
     GenerateResponseChunk,
     GenerationCommonConfig,
+    GenkitRegistry,
     Message,
     ModelInfo,
     Role,
     Supports,
-    TextPart,
     ToolDefinition,
 )
 from genkit.plugins.google_genai.models.utils import PartConverter
 
+
+class GeminiConfigSchema(genai_types.GenerateContentConfig):
+    pass
+
+
 GEMINI_1_0_PRO = ModelInfo(
     label='Google AI - Gemini Pro',
     versions=['gemini-pro', 'gemini-1.0-pro-latest', 'gemini-1.0-pro-001'],
@@ -605,6 +610,8 @@ def _genkit_to_googleai_cfg(
                     temperature=request_config.temperature,
                     stop_sequences=request_config.stop_sequences,
                 )
+            elif isinstance(request_config, GeminiConfigSchema):
+                cfg = request_config
             elif isinstance(request_config, dict):
                 cfg = genai_types.GenerateContentConfig(**request_config)
 

py/samples/hello-google-genai/src/hello.py

@@ -28,14 +28,15 @@
 from genkit.plugins.google_ai.models import gemini
 from genkit.plugins.google_genai import (
     EmbeddingTaskType,
+    GeminiConfigSchema,
     GeminiEmbeddingModels,
     GoogleGenai,
     google_genai_name,
 )
 
 ai = Genkit(
     plugins=[GoogleGenai()],
-    model=google_genai_name('gemini-2.0-flash'),
+    model=google_genai_name('gemini-2.0-flash-exp'),
 )
 
 
@@ -191,6 +192,15 @@ async def generate_character_unconstrained(name: str, ctx):
     return result.output
 
 
+@ai.flow()
+async def generate_images(name: str, ctx):
+    result = await ai.generate(
+        prompt=f'tell me a about the Eifel Tower with photos',
+        config=GeminiConfigSchema(response_modalities=['text', 'image']),
+    )
+    return result
+
+
 async def main() -> None:
     print(await say_hi(', tell me a joke'))