feat(py): Add type safety to Python SDK #4309

.github/workflows/python.yml

@@ -67,6 +67,11 @@ jobs:
           uv pip install -e .[dev,test,docs]
           uv pip install --group dev
 
+      - name: Install type checkers
+        run: |
+          cd py
+          uv pip install pyright mypy
+
       - name: Install NPM packages
         run: |
           npm install -g license-checker
@@ -96,6 +101,22 @@ jobs:
           fi
 
       - name: Type check with Ty
+        env:
+          PYTHONPATH: >-
+            packages/genkit/src:
+            plugins/anthropic/src:
+            plugins/compat-oai/src:
+            plugins/deepseek/src:
+            plugins/dev-local-vectorstore/src:
+            plugins/evaluators/src:
+            plugins/firebase/src:
+            plugins/flask/src:
+            plugins/google-cloud/src:
+            plugins/google-genai/src:
+            plugins/mcp/src:
+            plugins/ollama/src:
+            plugins/vertex-ai/src:
+            plugins/xai/src
         run: uv run --directory py ty check .
 
       - name: Run security scan

bin/check_license

@@ -45,6 +45,7 @@ $HOME/go/bin/addlicense \
   -ignore '**/.dist/**/*' \
   -ignore '**/.eggs/**/*' \
   -ignore '**/.idea/**/*' \
+  -ignore '**/.nox/**/*' \
   -ignore '**/.mypy_cache/**/*' \
   -ignore '**/.next/**/*' \
   -ignore '**/.output/**/*' \

py/bin/sanitize_schema_typing.py

@@ -40,6 +40,8 @@
 - Add docstrings if missing.
 """
 
+from __future__ import annotations
+
 import ast
 import sys
 from datetime import datetime
@@ -65,8 +67,10 @@ def is_rootmodel_class(self, node: ast.ClassDef) -> bool:
                     return True
         return False
 
-    def create_model_config(self, existing_config: ast.Call | None = None, frozen: bool = False) -> ast.Assign:
-        """Create or update a model_config assignment.
+    def create_model_config(self, existing_config: ast.Call | None = None, frozen: bool = False) -> ast.AnnAssign:
+        """Create or update a model_config assignment with proper type annotation.
+
+        Creates: model_config: ClassVar[ConfigDict] = ConfigDict(...)
 
         Ensures alias_generator=to_camel, populate_by_name=True, and extra='forbid',
         keeping other existing settings.
@@ -127,19 +131,31 @@ def create_model_config(self, existing_config: ast.Call | None = None, frozen: b
         # Sort keywords for consistent output (optional but good practice)
         keywords.sort(key=lambda kw: kw.arg or '')
 
-        return ast.Assign(
-            targets=[ast.Name(id='model_config')],
+        # Create ClassVar[ConfigDict] annotation
+        annotation = ast.Subscript(
+            value=ast.Name(id='ClassVar', ctx=ast.Load()),
+            slice=ast.Name(id='ConfigDict', ctx=ast.Load()),
+            ctx=ast.Load(),
+        )
+
+        return ast.AnnAssign(
+            target=ast.Name(id='model_config', ctx=ast.Store()),
+            annotation=annotation,
             value=ast.Call(func=ast.Name(id='ConfigDict'), args=[], keywords=keywords),
+            simple=1,
         )
 
-    def has_model_config(self, node: ast.ClassDef) -> ast.Assign | None:
+    def has_model_config(self, node: ast.ClassDef) -> ast.Assign | ast.AnnAssign | None:
         """Check if class already has model_config assignment and return it."""
         for item in node.body:
             if isinstance(item, ast.Assign):
                 targets = item.targets
                 if len(targets) == 1 and isinstance(targets[0], ast.Name):
                     if targets[0].id == 'model_config':
                         return item
+            elif isinstance(item, ast.AnnAssign):
+                if isinstance(item.target, ast.Name) and item.target.id == 'model_config':
+                    return item
         return None
 
     def visit_AnnAssign(self, node: ast.AnnAssign) -> ast.AnnAssign:  # noqa: N802
@@ -250,21 +266,39 @@ def visit_ClassDef(self, node: ast.ClassDef) -> object:  # noqa: N802
         if self.is_rootmodel_class(node):
             # Remove model_config from RootModel classes
             for stmt in node.body[body_start_index:]:
-                # Skip existing model_config
+                # Skip existing model_config (both Assign and AnnAssign)
                 if isinstance(stmt, ast.Assign) and any(
                     isinstance(target, ast.Name) and target.id == 'model_config' for target in stmt.targets
                 ):
                     self.modified = True  # Mark modified even if removing
                     continue
+                if (
+                    isinstance(stmt, ast.AnnAssign)
+                    and isinstance(stmt.target, ast.Name)
+                    and stmt.target.id == 'model_config'
+                ):
+                    self.modified = True
+                    continue
                 new_body.append(stmt)
         elif any(isinstance(base, ast.Name) and base.id == 'BaseModel' for base in node.bases):
             # Add or update model_config for BaseModel classes
             added_config = False
             frozen = node.name == 'PathMetadata'
             for stmt in node.body[body_start_index:]:
+                # Check for model_config (both Assign and AnnAssign)
+                is_model_config = False
                 if isinstance(stmt, ast.Assign) and any(
                     isinstance(target, ast.Name) and target.id == 'model_config' for target in stmt.targets
                 ):
+                    is_model_config = True
+                elif (
+                    isinstance(stmt, ast.AnnAssign)
+                    and isinstance(stmt.target, ast.Name)
+                    and stmt.target.id == 'model_config'
+                ):
+                    is_model_config = True
+
+                if is_model_config:
                     # Update existing model_config
                     updated_config = self.create_model_config(existing_model_config_call, frozen=frozen)
                     # Check if the config actually changed
@@ -354,6 +388,21 @@ def _inject_resources_field(self, body: list[ast.stmt | ast.Constant | ast.Assig
         self.modified = True
 
 
+def fix_field_defaults(content: str) -> str:
+    """Fix Field(None) and Field(None, ...) to use default=None for pyright compatibility.
+
+    Pyright doesn't recognize Field(None) as providing a default value,
+    but it does recognize Field(default=None).
+    """
+    import re
+
+    # Replace Field(None) with Field(default=None)
+    content = content.replace('Field(None)', 'Field(default=None)')
+    # Replace Field(None, other_args) with Field(default=None, other_args)
+    content = re.sub(r'Field\(None,', 'Field(default=None,', content)
+    return content
+
+
 def add_header(content: str) -> str:
     """Add the generated header to the content."""
     header = '''# Copyright {year} Google LLC
@@ -385,14 +434,11 @@ def add_header(content: str) -> str:
     # Ensure there's exactly one newline between header and content
     # and future import is right after the header block's closing quotes.
     future_import = 'from __future__ import annotations'
-    str_enum_block = """
+    compat_import_block = """
 import sys
+from typing import ClassVar
 
-if sys.version_info < (3, 11):
-    from strenum import StrEnum
-else:
-    from enum import StrEnum
-
+from genkit.core._compat import StrEnum
 from pydantic.alias_generators import to_camel
 """
 
@@ -405,7 +451,7 @@ def add_header(content: str) -> str:
     ]
     cleaned_content = '\n'.join(filtered_lines)
 
-    final_output = header_text + future_import + '\n' + str_enum_block + '\n\n' + cleaned_content
+    final_output = header_text + future_import + '\n' + compat_import_block + '\n\n' + cleaned_content
     if not final_output.endswith('\n'):
         final_output += '\n'
     return final_output
@@ -441,6 +487,9 @@ def process_file(filename: str) -> None:
         ast.fix_missing_locations(modified_tree)
         modified_source_no_header = ast.unparse(modified_tree)
 
+        # Fix Field(None) to Field(default=None) for pyright compatibility
+        modified_source_no_header = fix_field_defaults(modified_source_no_header)
+
         # Add header and specific imports correctly
         final_source = add_header(modified_source_no_header)
 

py/packages/genkit/README.md

@@ -18,7 +18,7 @@ pip install genkit-plugin-google-genai
 
 ```python
 from pydantic import BaseModel, Field
-from genkit.ai import Genkit
+from genkit.ai import Genkit, Output
 from genkit.plugins.google_genai import GoogleAI
 
 ai = Genkit(
@@ -34,10 +34,10 @@ class RpgCharacter(BaseModel):
     abilities: list[str] = Field(description='list of abilities (3-4)')
 
 @ai.flow()
-async def generate_character(name: str):
+async def generate_character(name: str) -> RpgCharacter:
     result = await ai.generate(
         prompt=f'generate an RPG character named {name}',
-        output_schema=RpgCharacter,
+        output=Output(schema=RpgCharacter),
     )
     return result.output
 

py/packages/genkit/pyproject.toml

@@ -46,6 +46,7 @@ dependencies = [
   "python-multipart>=0.0.22",
   "sse-starlette>=2.2.1",
   "pillow",
+  "typing_extensions>=4.0",
   "strenum>=0.4.15; python_version < '3.11'",
   "dotpromptz>=0.1.4",
   "uvicorn>=0.34.0",

py/packages/genkit/src/genkit/__init__.py

@@ -0,0 +1,119 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Genkit - Build AI-powered applications with ease.
+
+Genkit is an open-source Python toolkit designed to help you build
+AI-powered features in web and mobile apps.
+
+Basic usage:
+    from genkit import Genkit
+    from genkit.plugins.google_genai import GoogleAI
+
+    ai = Genkit(plugins=[GoogleAI()])
+
+    @ai.flow()
+    async def hello(name: str) -> str:
+        response = await ai.generate(model="gemini-2.0-flash", prompt=f"Hello {name}")
+        return response.text
+"""
+
+# Main class
+# Re-export everything from genkit.ai for backwards compatibility
+from genkit.ai import (
+    GENKIT_CLIENT_HEADER,
+    GENKIT_VERSION,
+    ActionKind,
+    ActionRunContext,
+    Chat,
+    ChatOptions,
+    ChatStreamResponse,
+    ExecutablePrompt,
+    FlowWrapper,
+    GenerateStreamResponse,
+    GenkitRegistry,
+    OutputOptions,
+    PromptGenerateOptions,
+    ResumeOptions,
+    SimpleRetrieverOptions,
+    ToolRunContext,
+    tool_response,
+)
+from genkit.ai._aio import Genkit, Output
+
+# Core types for convenience (also available from genkit.types)
+from genkit.blocks.document import Document
+
+# Response types
+from genkit.blocks.model import GenerateResponseWrapper
+
+# Setup plugin discovery (must be done before any plugin imports)
+from genkit.core._plugins import extend_plugin_namespace
+
+# Errors (user-facing)
+from genkit.core.error import GenkitError, UserFacingError
+
+# Plugin interface
+from genkit.core.plugin import Plugin
+from genkit.core.typing import (
+    Media,
+    MediaPart,
+    Message,
+    Part,
+    Role,
+    TextPart,
+)
+
+extend_plugin_namespace()
+
+__all__ = [
+    # Main class
+    'Genkit',
+    'Output',
+    # Response types
+    'GenerateResponseWrapper',
+    # Errors
+    'GenkitError',
+    'UserFacingError',
+    # Core types (convenience)
+    'Document',
+    'Message',
+    'Part',
+    'Role',
+    'TextPart',
+    'MediaPart',
+    'Media',
+    # Plugin interface
+    'Plugin',
+    # From genkit.ai
+    'ActionKind',
+    'ActionRunContext',
+    'Chat',
+    'ChatOptions',
+    'ChatStreamResponse',
+    'ExecutablePrompt',
+    'FlowWrapper',
+    'GenerateStreamResponse',
+    'GenkitRegistry',
+    'OutputOptions',
+    'PromptGenerateOptions',
+    'ResumeOptions',
+    'SimpleRetrieverOptions',
+    'ToolRunContext',
+    'tool_response',
+    'GENKIT_CLIENT_HEADER',
+    'GENKIT_VERSION',
+]